Speedup Segformer decoder (#998)

qubvel-org · Dec 6, 2024 · 5a1a8f1 · 5a1a8f1
1 parent 05b48d5
commit 5a1a8f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/segmentation_models_pytorch/decoders/segformer/decoder.py b/segmentation_models_pytorch/decoders/segformer/decoder.py
@@ -15,7 +15,7 @@ def forward(self, x: torch.Tensor):
         batch, _, height, width = x.shape
         x = x.flatten(2).transpose(1, 2)
         x = self.linear(x)
-        x = x.transpose(1, 2).reshape(batch, -1, height, width).contiguous()
+        x = x.transpose(1, 2).reshape(batch, -1, height, width)
         return x