Here is code snippet I’m using and this is the input part of the model:
  

      github.com
  

  
    bshall/ZeroSpeech/blob/0ef25e4daac07888d22b802c649480dc70b719fe/model.py#L20



      
      
          
          
def get_gru_cell(gru):
              gru_cell = nn.GRUCell(gru.input_size, gru.hidden_size)
              gru_cell.weight_hh.data = gru.weight_hh_l0.data
              gru_cell.weight_ih.data = gru.weight_ih_l0.data
              gru_cell.bias_hh.data = gru.bias_hh_l0.data
              gru_cell.bias_ih.data = gru.bias_ih_l0.data
              return gru_cell
          
          

          
class Encoder(nn.Module):
              def __init__(self, in_channels, channels, n_embeddings, embedding_dim, jitter=0):
                  super(Encoder, self).__init__()
                  self.encoder = nn.Sequential(
                      nn.Conv1d(in_channels, channels, 3, 1, 0, bias=False),
                      nn.BatchNorm1d(channels),
                      nn.ReLU(True),
                      nn.Conv1d(channels, channels, 3, 1, 1, bias=False),
                      nn.BatchNorm1d(channels),
                      nn.ReLU(True),
                      nn.Conv1d(channels, channels, 4, 2, 1, bias=False),
      
    
  


  

  
    
    
  

  


The argument values which have been passed to this object is:
in_channels: 80
channels: 768
n_embeddings: 512
embedding_dim: 64
jitter: 0.5
The models are initialised here:
  

      github.com
  

  
    bshall/ZeroSpeech/blob/0ef25e4daac07888d22b802c649480dc70b719fe/train.py#L39



      
      
              print("Saved checkpoint: {}".format(checkpoint_path.stem))
          
          

          
@hydra.main(config_path="config/train.yaml")
          def train_model(cfg):
              tensorboard_path = Path(utils.to_absolute_path("tensorboard")) / cfg.checkpoint_dir
              checkpoint_dir = Path(utils.to_absolute_path(cfg.checkpoint_dir))
              writer = SummaryWriter(tensorboard_path)
              device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
          
          
    encoder = Encoder(**cfg.model.encoder)
              decoder = Decoder(**cfg.model.decoder)
              encoder.to(device)
              decoder.to(device)
          
          
    optimizer = optim.Adam(
                  chain(encoder.parameters(), decoder.parameters()),
                  lr=cfg.training.optimizer.lr)
              [encoder, decoder], optimizer = amp.initialize([encoder, decoder], optimizer, opt_level="O1")
              scheduler = optim.lr_scheduler.MultiStepLR(
                  optimizer, milestones=cfg.training.scheduler.milestones,
      
    
  


  

  
    
    
  

  


Let me know if anything you need to assure.
By the way, this is a appreciation for you @ptrblck. Whenever I face any pytorch issues and look for the solution, it’s often you who comes with the best solution. Thanks for everything bud!