Fix detachment again

libs/sepconv/SeparableConvolution.py

@@ -12,9 +12,10 @@ def __init__(self):
         super(SeparableConvolution, self).__init__()
     # end
 
-    def forward(self, input, vertical, horizontal):
+    @staticmethod
+    def forward(context, input, vertical, horizontal):
 
-        self.save_for_backward(input, vertical, horizontal)
+        context.save_for_backward(input, vertical, horizontal)
 
         intBatches = input.size(0)
         intInputDepth = input.size(1)
@@ -50,9 +51,10 @@ def forward(self, input, vertical, horizontal):
         return output
     # end
 
-    def backward(self, grad_output):
+    @staticmethod
+    def backward(context, grad_output):
 
-        _input, vertical, horizontal = self.saved_tensors
+        _input, vertical, horizontal = context.saved_tensors
 
         grad_input = _input.new().resize_(_input.size()).zero_()
         grad_vertical = vertical.new().resize_(vertical.size()).zero_()

src/main.py

@@ -75,8 +75,6 @@ def train(epoch):
         print('Gradients ready.')
         optimizer.step()
 
-        detach_all(model.parameters())
-
         print("===> Epoch[{}]({}/{}): Loss: {:.4f}".format(epoch, iteration, len(training_data_loader), loss.item()))
     print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, epoch_loss / len(training_data_loader)))
 

src/model.py

@@ -5,14 +5,13 @@
 import torch
 import torch.nn as nn
 import torch.nn.init as init
+from torch.autograd import Variable, gradcheck
 from torch.nn.modules.loss import _Loss, _assert_no_grad
 from src.separable_convolution import SeparableConvolutionSlow
 from libs.sepconv.SeparableConvolution import SeparableConvolution
 import src.config as config
 import src.interpolate as interpol
-import numpy as np
 
-from torch.autograd import Variable, gradcheck
 
 class Net(nn.Module):
 
@@ -50,7 +49,7 @@ def __init__(self):
         self.pad = nn.ConstantPad2d(sep_kernel // 2, 0.0)
 
         if torch.cuda.is_available() and not config.ALWAYS_SLOW_SEP_CONV:
-            self.separable_conv = SeparableConvolution()
+            self.separable_conv = SeparableConvolution.apply
         else:
             self.separable_conv = SeparableConvolutionSlow()