Removed MinkowskiConvolutionFunction from MinkowskiConvolution to all…

…ow pickling (#139)
NVIDIA · May 15, 2020 · 6618442 · 6618442
1 parent c351d10
commit 6618442
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/MinkowskiEngine/MinkowskiConvolution.py b/MinkowskiEngine/MinkowskiConvolution.py
@@ -268,9 +268,13 @@ def forward(self,
             outfeat = input.F.mm(self.kernel)
             out_coords_key = input.coords_key
         else:
+            if self.is_transpose:
+                conv = MinkowskiConvolutionTransposeFunction()
+            else:
+                conv = MinkowskiConvolutionFunction()
             # Get a new coords key or extract one from the coords
             out_coords_key = _get_coords_key(input, coords)
-            outfeat = self.conv.apply(input.F, self.kernel, input.tensor_stride,
+            outfeat = conv.apply(input.F, self.kernel, input.tensor_stride,
                                       self.stride, self.kernel_size,
                                       self.dilation, self.region_type_,
                                       self.region_offset_, input.coords_key,
@@ -384,7 +388,6 @@ def __init__(self,
             is_transpose=False,
             dimension=dimension)
         self.reset_parameters()
-        self.conv = MinkowskiConvolutionFunction()
 
 
 class MinkowskiConvolutionTranspose(MinkowskiConvolutionBase):
@@ -458,7 +461,6 @@ def __init__(self,
             dimension=dimension)
         self.reset_parameters(True)
         self.generate_new_coords = generate_new_coords
-        self.conv = MinkowskiConvolutionTransposeFunction()
 
     def forward(self,
                 input: SparseTensor,
@@ -487,7 +489,7 @@ def forward(self,
         else:
             # Get a new coords key or extract one from the coords
             out_coords_key = _get_coords_key(input, coords, tensor_stride=1)
-            outfeat = self.conv.apply(
+            outfeat = MinkowskiConvolutionTransposeFunction().apply(
                 input.F, self.kernel, input.tensor_stride, self.stride,
                 self.kernel_size, self.dilation, self.region_type_,
                 self.region_offset_, self.generate_new_coords, input.coords_key,
@@ -496,4 +498,4 @@ def forward(self,
             outfeat += self.bias
 
         return SparseTensor(
-            outfeat, coords_key=out_coords_key, coords_manager=input.coords_man)
+            outfeat, coords_key=out_coords_key, coords_manager=input.coords_man)