DOFA: fix bug in patch embedding

microsoft · Aug 27, 2024 · e5cbb06 · e5cbb06
1 parent 27a981c
commit e5cbb06
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchgeo/models/dofa.py b/torchgeo/models/dofa.py
@@ -203,8 +203,10 @@ def forward(self, x: Tensor, wavelengths: Tensor) -> tuple[Tensor, Tensor]:
         weight, bias = self.weight_generator(waves)  # 3x3x3
 
         dynamic_weight = weight.view(
-            self.embed_dim, inplanes, self.kernel_size, self.kernel_size
-        )  # 3xoutdx16x16
+            inplanes, self.kernel_size, self.kernel_size, self.embed_dim
+        )
+        dynamic_weight = dynamic_weight.permute([3, 0, 1, 2])
+
         if bias is not None:
             bias = bias.view([self.embed_dim]) * self.scaler