Hi @jacob,
I implemented a couple of custom RNNs recently by simply modifying this code,

  
      github.com
  
  
    jihunchoi/recurrent-batch-normalization-pytorch/blob/master/bnlstm.py 903
"""Implementation of batch-normalized LSTM."""
import torch
from torch import nn
from torch.autograd import Variable
from torch.nn import functional, init


class SeparatedBatchNorm1d(nn.Module):

    """
    A batch normalization module which keeps its running mean
    and variance separately per timestep.
    """

    def __init__(self, num_features, max_length, eps=1e-5, momentum=0.1,
                 affine=True):
        """
        Most parts are copied from
        torch.nn.modules.batchnorm._BatchNorm.
        """


  This file has been truncated. show original

  
  
    
    
  
  


Hope that helps,
Best,
Ajay