I’m trying to implement Variational Dropout to Recurrent Highway Network.
Because my implementation of the RHN is in for loop for timesteps, I need to save the mask for every time-step for future use. In each batch, I’m creating a new mask.
This is my function for creating the mask:
class MaskDropout(nn.Module):
    """Same as Lockdropout, but for a single time-step - using the same mask"""
    def __init__(self, dropout=0.3):
        super(MaskDropout, self).__init__()
        self.mask = None
        self.dropout = dropout

    def reset(self, x):
        """ MUST BE CALLED BEFORE FORWARD! call this function every batch to reset the mask """
        m = x.data.new(x.size()).bernoulli_(1 - self.dropout)
        m = m.div_(1 - self.dropout)
        self.mask = torch.autograd.Variable(m, requires_grad=False)

    def forward(self, x):
        """forward function for dropout implementation (mask is same)"""
        if not self.training or not self.dropout:
            return x
        return self.mask * x

This is an example of feedforward with a mask:
    def forward(self, input, s_t0_l0):
        _st_batch_size = s_t0_l0.size(0)
        _st_hid_dim = s_t0_l0.size(1)

        _timesteps = input.size(0)
        _batch_size = input.size(1)
        _features = input.size(2)

        inputs = list(input.unbind(0))  # splitting inputs into time-steps (list of ts elements of [bs x feats])
        s_t_minus_1 = s_t0_l0

        self.maskdrop.reset(s_t0_l0)                # resetting the variational dropout masks

        for t in range(len(inputs)):
            # s_t_L = output of RHN only
            s_t_L = self.rhncell(inputs[t], s_t_minus_1)  # RHN Cell
            # s_t = output of HSG (after RHN) - This is an extension, you can skip this line..
            s_t = self.hsgcell(s_t_L, s_t_minus_1)  # HSG Cell

            # variational dropout for s_t
            s_t = self.maskdrop(s_t)     

            s_t_minus_1 = s_t
            outputs += [s_t]

        rhn_outputs = torch.stack(outputs)
        dense_out = self.output_fc(rhn_outputs)
        ...

Now the problem is that with the following line, m = m.div_(1 - self.dropout) , at some point, there are nan values. Without it everything is perfect, but it’s a wrong implementation of dropout in general…
My theory is that multiplying each time the same values by a value greater than 1 will cause gradients explosion, but this has been done in the past (as you can see from implementing variational dropout for LSTM networks…)
What am I missing here?
Thank you!