Here’s what I am trying to do:
class StyleLoss(nn.Module):

    def __init__(self, strength, normalize):
        super(StyleLoss, self).__init__()
        #self.target = torch.Tensor()
        self.target = Variable((torch.Tensor()),requires_grad=False)
        self.strength = strength
        self.gram = GramMatrix()
        self.criterion = nn.MSELoss()
        self.mode = None
        self.blend_weight = None
        self.G = None
        self.normalize = 'False'


    def forward(self, input):
        self.output = input.clone()
        self.G = self.gram(input)
        self.G.div(input.nelement())
        if self.mode == 'capture':
          if self.blend_weight == None:
            self.target.data.resize_as_(self.G.cpu().data).copy_(self.G.cpu().data)
          elif self.target.nelement() == 0:
            self.target.data.resize_as_(self.G.cpu().data).copy_(self.G.cpu().data).mul_(self.blend_weight.cpu())
          else:
            self.target.data.add(self.blend_weight.cpu(), self.G.data)
        elif self.mode == 'loss':
            self.loss = self.strength * self.criterion(self.G.cpu(), self.target.cpu())
        return self.output

After running the forward function, I need to run self.G backwards to create dG, like in this Lua/Torch code: dG = self.crit:backward(self.G, self.target), but unlike Lua/Torch, PyTorch doesn’t have a backwards pass function (at least one that works in the same way) like Lua/Torch.
    def backward(self, input, gradOutput):
        if self.mode == 'loss':
          dG = self.G.backward(self.G, self.target)
          self.gradInput = self.gram.backward(input, self.dG)
          if self.normalize == 'True':
	         self.gradInput.div(torch.norm(self.gradInput, 1) + 1e-8) # Normalize Gradients
	      self.gradInput.mul(self.strength)
          self.gradInput.add(gradOutput)
        else:
          self.gradInput = gradOutput
        return self.gradInput

Could someone please help me with that this? Should I somehow be using a backwards hook? If so, how would that work?