You just take an LSTMCell implementation and modify it.
For example, you can copy the implementation here and modify it for your own purposes:

  
      github.com
  
  
    pytorch/benchmark/blob/09eaadc1d05ad442b1f0beb82babf875bbafb24b/rnns/fastrnns/cells.py#L25-L40 45

def lstm_cell(input, hidden, w_ih, w_hh, b_ih, b_hh):
# type: (Tensor, Tuple[Tensor, Tensor], Tensor, Tensor, Tensor, Tensor) -> Tuple[Tensor, Tensor]
hx, cx = hidden
gates = torch.mm(input, w_ih.t()) + torch.mm(hx, w_hh.t()) + b_ih + b_hh


ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)


ingate = torch.sigmoid(ingate)
forgetgate = torch.sigmoid(forgetgate)
cellgate = torch.tanh(cellgate)
outgate = torch.sigmoid(outgate)


cy = (forgetgate * cx) + (ingate * cellgate)
hy = outgate * torch.tanh(cy)


return hy, cy