This is the model definition. The sizes are below the definition. The model is made up of 2 VAEs. The input inp_pose and key_pose are each torch.Size([128, 63]).
class model(nn.Module):
    def __init__(self, args):
        super(model, self).__init__()
                 
        for key, value in args.__dict__.items():
            setattr(self, key, value)
                        
        """
        Pose
        """
          
        self.inp_pose_encoder = make_mlp([self.pose_encoder_units[0]+3]+self.pose_encoder_units[1:],self.pose_encoder_activations)
        self.key_pose_encoder = make_mlp(self.pose_encoder_units,self.pose_encoder_activations)
        self.pose_mu          = make_mlp(self.pose_mu_var_units,self.pose_mu_var_activations)
        self.pose_log_var     = make_mlp(self.pose_mu_var_units,self.pose_mu_var_activations)
        self.key_pose_decoder = make_mlp(self.pose_decoder_units,self.pose_decoder_activations)
        
        """
        Time
        """
        
        self.delta_pose_encoder = make_mlp(self.delta_pose_encoder_units,self.delta_pose_encoder_activations)
        self.time_encoder       = make_mlp(self.time_encoder_units,self.time_encoder_activations)
        self.time_mu            = make_mlp(self.time_mu_var_units,self.time_mu_var_activations)
        self.time_log_var       = make_mlp(self.time_mu_var_units,self.time_mu_var_activations)
        self.time_decoder       = make_mlp(self.time_decoder_units,self.time_decoder_activations)
        
        self.norm = tdist.Normal(torch.tensor(0.0), torch.tensor(1.0))
                        
    def forward(self, data, mode):
    
        #batch_size = data["inp_pose"].shape[0]
    
        inp_pose = data["inp_pose"].view(self.batch_size,-1) 
        key_pose = data["key_pose"].view(self.batch_size,-1)
        key_object = data["key_object"].view(self.batch_size,-1)
               
        """
        compute pose
        """
    
        # feed x and y
        inp_pose_features = torch.cat((inp_pose, key_object), dim=1)
        inp_pose_features = self.inp_pose_encoder(inp_pose_features)
        key_pose_features = self.key_pose_encoder(key_pose)
        
        # get gaussian parameters
        pose_posterior = torch.cat((inp_pose_features,key_pose_features),dim=1)
        pose_posterior_mu = self.pose_mu(pose_posterior)
        pose_posterior_log_var = self.pose_log_var(pose_posterior)
        
        # sample
        pose_posterior_std = torch.exp(0.5*pose_posterior_log_var)
        pose_posterior_eps = self.norm.sample([self.batch_size, pose_posterior_mu.shape[1]]).cuda()
        pose_posterior_z   = pose_posterior_mu + pose_posterior_eps*pose_posterior_std
        
        z_p = pose_posterior_z if mode == "tr" else self.norm.sample([self.batch_size, self.pose_mu_var_units[-1]]).cuda()
        
        # forecast
        pred_key_pose = torch.cat((z_p,inp_pose_features),dim=1)
        pred_key_pose = self.key_pose_decoder(pred_key_pose)
        
        """
        compute time
        """
        
        # compute delta_pose
        delta_pose = key_pose - inp_pose
        
        # feed x and y
        delta_pose_features = self.delta_pose_encoder(delta_pose)
        time_features = self.time_encoder(data["time"].unsqueeze(1))
        
        # get gaussian parameters
        time_posterior = torch.cat((delta_pose_features,time_features),dim=1)
        time_posterior_mu = self.time_mu(time_posterior)
        time_posterior_log_var = self.time_log_var(time_posterior)
        
        # sample
        time_posterior_std = torch.exp(0.5*time_posterior_log_var)
        time_posterior_eps = self.norm.sample([self.batch_size, time_posterior_mu.shape[1]]).cuda()
        time_posterior_z   = time_posterior_mu + time_posterior_eps*time_posterior_std
        
        z_t = time_posterior_z if mode == "tr" else self.norm.sample([self.batch_size, self.time_mu_var_units[-1]]).cuda()
              
        # compute time
        time = torch.cat((z_t,delta_pose_features),dim=1)
        time = torch.squeeze(self.time_decoder(time))   
        
        return {"key_pose":pred_key_pose.view(self.batch_size,21,3), "time":time,
                "pose_posterior":{"mu":pose_posterior_mu, "log_var":pose_posterior_log_var}, 
                "time_posterior":{"mu":time_posterior_mu, "log_var":time_posterior_log_var}}

And the specific units
model(
  (inp_pose_encoder): Sequential(
    (0): Linear(in_features=66, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=128, bias=True)
    (3): ReLU()
  )
  (key_pose_encoder): Sequential(
    (0): Linear(in_features=63, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=128, bias=True)
    (3): ReLU()
  )
  (pose_mu): Sequential(
    (0): Linear(in_features=256, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=8, bias=True)
  )
  (pose_log_var): Sequential(
    (0): Linear(in_features=256, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=8, bias=True)
  )
  (key_pose_decoder): Sequential(
    (0): Linear(in_features=136, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=63, bias=True)
  )
  (delta_pose_encoder): Sequential(
    (0): Linear(in_features=63, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=128, bias=True)
    (3): ReLU()
  )
  (time_encoder): Sequential(
    (0): Linear(in_features=1, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=16, bias=True)
    (3): ReLU()
  )
  (time_mu): Sequential(
    (0): Linear(in_features=144, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=8, bias=True)
  )
  (time_log_var): Sequential(
    (0): Linear(in_features=144, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=8, bias=True)
  )
  (time_decoder): Sequential(
    (0): Linear(in_features=136, out_features=256, bias=True)
    (1): ReLU()
    (2): Linear(in_features=256, out_features=1, bias=True)
    (3): ReLU()
  )
)