Update codes.py

codes.py

@@ -104,12 +104,12 @@ def mtrx3(Affine_mtrx):
     mtrx_shape = Affine_mtrx.shape
     if len(mtrx_shape)==3:
         N_Mbatches = mtrx_shape[0]
-        AM3 = torch.zeros( [N_Mbatches,3,3]).to(device)
+        AM3 = torch.zeros( [N_Mbatches,3,3])#.to(device)
         AM3[:,0:2,:] = Affine_mtrx
         AM3[:,2,2] = 1
     elif len(mtrx_shape)==2:
         N_Mbatches = 1
-        AM3 = torch.zeros([3,3]).to(device)
+        AM3 = torch.zeros([3,3])#.to(device)
         AM3[0:2,:] = Affine_mtrx
         AM3[2,2] = 1 
     return AM3