Here is the full along with the imports
import torch
import torch.nn as nn
import torch.nn.functional as F

class VariationalLSTM(nn.Module):
    
    
    
    """
    Modified by Code taken at  https://github.com/fastai/fastai/blob/master/fastai/text/models/awd_lstm.py#L27.
    FastAI implementation implements WeightDropLSTM which is different from VariationaLSTM.
    Former drops certain elements in matrix that defines recurrent connection in each layer
    Variational Dropout(VDropout) is different in following.Hence those changes are made
    
    
    i)VDropout Should zero out rows of matrices, instead of arbitrarily zeroing it out 
    ii)VDropout also zeros matrices that involves Output connection 
    
    """    


    def __init__(self,module,weight_p=0.0, layer_names=['weight_hh_l0']):
        
        super().__init__()
        
        self.module  = module
        self.weight_p,self.layer_names = weight_p,layer_names
        self.idxs = [] if hasattr(self.module, '_flat_weights_names') else None
        
        
        
        for layer in self.layer_names:
        
        
            w = getattr(self.module, layer)
            self.register_parameter(f'{layer}_raw', nn.Parameter(w.data))
              """
              following is the modification I made
             """   
            
            N,K  = w.shape
            mask = F.dropout(torch.ones(N,1),p=self.weight_p,training=False)
            mask = mask.repeat(1,K)
            self.module._parameters[layer] = w * mask
            
            if self.idxs is not None: self.idxs.append(self.module._flat_weights_names.index(layer))
        if isinstance(self.module, (nn.RNNBase, nn.modules.rnn.RNNBase)):
            self.module.flatten_parameters = self._do_nothing

    def _setweights(self):
        for i,layer in enumerate(self.layer_names):
            raw_w = getattr(self, f'{layer}_raw')


           """
              following is the modification I made
          """   
            N,K  = raw_w.shape
            mask = F.dropout(torch.ones(N,1),p=self.weight_p)
            mask = mask.repeat(1,K)
            self.module._parameters[layer] = raw_w * mask
            
            if self.idxs is not None: self.module._flat_weights[self.idxs[i]] = self.module._parameters[layer]

    def forward(self, *args, do_dropout=True, **kwargs):
        
        if(do_dropout):
            self._setweights()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return self.module.forward(*args)
        
    def reset(self):
        for layer in self.layer_names:
            raw_w = getattr(self, f'{layer}_raw')
            self.module._parameters[layer] = F.dropout(raw_w, p=self.weight_p, training=False)
        if hasattr(self.module, 'reset'): self.module.reset()    
    
    def _do_nothing(self): pass

module = nn.LSTM(1,20)
m = VariationalLSTM(module , weight_p=0.25)
print(m)
optimizer = torch.optim.Adam(m.parameters() , lr = 1e-3)