Tkurth/distributed contig fix (#12)

* fixing contig in distributed sht

* make it normally contiguous before all to all

* commented the code to be a WAR

torch_harmonics/distributed/primitives.py

@@ -72,16 +72,22 @@ class distributed_transpose_azimuth(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, x, dim):
+        input_format = get_memory_format(x)
+        # WAR for a potential contig check torch bug for channels last contig tensors
+        x = x.contiguous()
         xlist, _ = _transpose(x, dim[0], dim[1], group=azimuth_group())
-        x = torch.cat(xlist, dim=dim[1])
+        x = torch.cat(xlist, dim=dim[1]).contiguous(memory_format=input_format)
         ctx.dim = dim
         return x
 
     @staticmethod
     def backward(ctx, go):
+        input_format = get_memory_format(go)
         dim = ctx.dim
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        go = go.contiguous()
         gilist, _ = _transpose(go, dim[1], dim[0], group=azimuth_group())
-        gi = torch.cat(gilist, dim=dim[0])
+        gi = torch.cat(gilist, dim=dim[0]).contiguous(memory_format=input_format)
         return gi, None
 
     
@@ -89,15 +95,21 @@ class distributed_transpose_polar(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, x, dim):
+        input_format = get_memory_format(x)
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        x = x.contiguous()
         xlist, _ = _transpose(x, dim[0], dim[1], group=polar_group())
-        x = torch.cat(xlist, dim=dim[1])
+        x = torch.cat(xlist, dim=dim[1]).contiguous(memory_format=input_format)
         ctx.dim = dim
         return x
 
     @staticmethod
     def backward(ctx, go):
+        input_format = get_memory_format(go)
         dim = ctx.dim
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        go = go.contiguous()
         gilist, _ = _transpose(go, dim[1], dim[0], group=polar_group())
-        gi = torch.cat(gilist, dim=dim[0])
+        gi = torch.cat(gilist, dim=dim[0]).contiguous(memory_format=input_format)
         return gi, None