Solved by dfalbel in post #4 
              
              
                I think scripting works fine in that case and is probably the best option: 
import torch
import torch.nn as nn

class MyGRU (nn.Module):
  def __init__ (self):
    super(MyGRU, self).__init__()
    self.gru = nn.GRU(10, 32, 5)
  
  def forward (self, x):
    return self.gru(x)

module = torch.jit.sc…