fixing: transformer forward method

Jorgedavyd · Jul 10, 2024 · 6fbef5b · 6fbef5b
1 parent d3ac3f6
commit 6fbef5b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/lightorch/nn/transformer/transformer.py b/lightorch/nn/transformer/transformer.py
@@ -78,6 +78,7 @@ def __init__(
         fc: Optional[nn.Module] = None,
         n_layers: int = 1,
     ) -> None:
+        assert (encoder is not None or decoder is not None), "Not valid parameters, must be at least one encoder or decoder."
         super().__init__()
         self.embedding = embedding_layer
         self.pe = positional_encoding
@@ -110,8 +111,8 @@ def forward(self, x: Tensor) -> Tensor:
         else:
             for decoder in self.decoder:
                 out = decoder(x)
-
-        x = self.fc(out)
+        if self.fc:
+            x = self.fc(out)
 
         return x