ModelGRU.py

class ModelGRU(nn.Module):
    def __init__(self):
        super().__init__()
        self.i_h = nn.Embedding(nv,nh)
        self.rnn = nn.GRU(nh, nh, 2, batch_first=True)
        self.h_o = nn.Linear(nh,nv)
        self.bn1 = nn.BatchNorm1d(num_features=nh)
        
    def forward(self, x):
      self.h = torch.zeros(2, x.shape[0], nh).cuda()
      res,h = self.rnn(self.i_h(x), self.h)
      self.h = h.detach()
      res = res.permute(0,2,1)
      res = self.bn1(res)
      return self.h_o(res.permute(0,2,1))