Upload latent_bayesian.py

latent_bayesian.py

@@ -197,7 +197,8 @@ def bayesian_forecast(in_tensor, n_channels, physio_channels,
     if len(action_channels) > 0:
         with torch.no_grad():
             future_action_enc_out, _, _ = latent_encoder.forward_encoder(in_tensor.clone()[:, action_channels, :], masking=False)
-            future_action_enc_out = future_action_enc_out.permute(1, 2, 0).flatten(start_dim=1)[1:1+(pred_length//latent_encoder.patch_size)+1, :] # sample_steps, embed_dim*bn*action_nvar
+            context_npatchs = context_length // latent_encoder.patch_size
+            future_action_enc_out = future_action_enc_out.permute(1, 2, 0).flatten(start_dim=1)[context_npatchs:context_npatchs+(pred_length//latent_encoder.patch_size)+1, :] # sample_steps, embed_dim*bn*action_nvar
     appended_embeds = quantile_traj_of_state(
         # enc_out[0, -1, :], 
         enc_out[-1:, :],