Merge pull request #52 from kaseris/fix/pvred-model-input

Fix/pvred model input
kaseris · Dec 14, 2023 · cdd1abe · cdd1abe
2 parents 3dcfb33 + 59d5277
commit cdd1abe
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/configs/pvred.yaml → configs/pvred_base.yaml b/configs/pvred.yaml → configs/pvred_base.yaml
@@ -14,9 +14,9 @@ transforms:
  - name: MinMaxScaleTransform
  args:
  feature_scale: [0.0, 1.0]
- - name: CartToExpMapsTransform
- args:
- parents: null
+ # - name: CartToExpMapsTransform
+ #  args:
+ #  parents: null
 
 loss:
  name: MSELoss
@@ -48,7 +48,7 @@ model:
  enc_type: lstm
  dec_type: lstm
  include_velocity: false
- pos_enc: add
+ pos_enc: null
  batch_first: true
  std_thresh: 0.0001
  use_std_mask: false

diff --git a/src/skelcast/models/rnn/pvred.py b/src/skelcast/models/rnn/pvred.py
@@ -182,7 +182,8 @@ def forward(self, x: torch.Tensor, y:torch.Tensor, masks: torch.Tensor = None) -
 
  batch_size, seq_len, n_bodies, n_joints, dims = x.shape
  x = x.view(batch_size, seq_len, n_bodies * n_joints * dims)
- masks = masks.view(batch_size, seq_len, n_bodies * n_joints * dims)
+ if masks is not None:
+ masks = masks.view(batch_size, seq_len, n_bodies * n_joints * dims)
  # Calculate the velocity if the include_velocity flag is true
  if self.include_velocity:
  vel_inp = self._calculate_velocity(x)