dataset/view bugfix

2023-01-14 18:42:27 +00:00
parent 581a5b2af2
commit 7dffb02286
2 changed files with 6 additions and 3 deletions
@@ -13,6 +13,7 @@ class ViewEncoder(nn.Module):
        super().__init__()
        self.features: int = features

+        self.flatten = nn.Flatten()
        self.lin1 = nn.LazyLinear(features).to(device)
        self.act1 = nn.GELU().to(device)
        self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
@@ -20,7 +21,8 @@ class ViewEncoder(nn.Module):
    

    def forward(self, input: torch.Tensor) -> torch.Tensor:
-        x: torch.Tensor = self.lin1(input)
+        x: torch.Tensor = self.flatten(input)
+        x= self.lin1(x)
        x = self.act1(x)
        x = self.lin2(x)
        x = self.act2(x)