A part of my arcitecture is trying to learn to weigh a multinomial distribution.
Below I added the forward function of the layer.
My problem is that for some reason the backward flow doesn’t access this layer’s backward() method so I cannot weigh the gradients as they should be.
I know I can use register_hook to manually fix the gradients of a certain layer but it won’t change the gradients of all the layers that came before it.
Any advice will be appreciated.
Forward layer:
def forward(self, input):
    self.output = Variable(torch.zeros(input.size()))
    self._index = torch.multinomial(input + constants.epsilon, 1, False)
    self.output.scatter_(1, self._index, torch.unsqueeze(self.one.repeat(self._index.size()[0]),1))
    return self._index.float()