small

app.py

@@ -371,7 +371,7 @@ def process_single_video(
         target_progress = np.linspace(0.0, 1.0, num=num_frames).tolist()
         success_label = [1.0 if prog > 0.5 else 0.0 for prog in target_progress]
 
-        # predict_last_frame_mask: server expects a list (1.0 per frame); omit and server pad_list_to_max can get None
+        # predict_last_frame_mask: server collator requires a list (1.0 per frame = no masking for inference)
         predict_last_frame_mask = [1.0] * num_frames
 
         # Create Trajectory
@@ -494,7 +494,7 @@ def process_two_videos(
         success_label_a = [1.0 if prog > 0.5 else 0.0 for prog in target_progress_a]
         success_label_b = [1.0 if prog > 0.5 else 0.0 for prog in target_progress_b]
 
-        # predict_last_frame_mask: server expects a list (1.0 per frame); None causes pad_list_to_max to fail
+        # predict_last_frame_mask: server collator requires a list per trajectory (1.0 = no masking)
         mask_a = [1.0] * num_frames_a
         mask_b = [1.0] * num_frames_b
 

dataset_types.py

@@ -35,6 +35,7 @@ class Trajectory(BaseModel):
     target_progress: Optional[Union[List[float], List[torch.Tensor], torch.Tensor, None]] = None
     partial_success: Optional[Union[float, torch.Tensor]] = None  # float for continuous, Tensor for C51 discrete
     success_label: Optional[List[float]] = None  # Success labels for each frame (1.0 for success, 0.0 for failure)
+    predict_last_frame_mask: Optional[List[float]] = None  # 1.0 per frame for inference (no masking); server requires a list
     metadata: Optional[Dict[str, Any]] = None
     data_gen_strategy: Optional[str] = None