Fix Ego4D target compatibility in local baselines

train_turn_head.py

@@ -224,6 +224,7 @@ def _load_examples(
     local_feature_root: Path | None = None,
     local_modalities: str = "av",
     local_contract: dict[str, Any] | None = None,
+    compute_transition_targets: bool = True,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -259,11 +260,19 @@ def _load_examples(
             if data.local_features is not None
             else torch.empty(len(states), 0)
         )
-        activity_sequence = torch.cat(
-            (data.targets[current_indices, None, :], targets),
-            dim=1,
-        )
-        transition_targets = transition_targets_from_activity(activity_sequence)
+        if compute_transition_targets:
+            activity_sequence = torch.cat(
+                (data.targets[current_indices, None, :], targets),
+                dim=1,
+            )
+            transition_targets = transition_targets_from_activity(
+                activity_sequence
+            )
+        else:
+            transition_targets = torch.zeros(
+                targets.shape[:-1],
+                dtype=targets.dtype,
+            )
         dense_example_count = len(states)
         if selected_indices is not None:
             if name not in selected_indices:
@@ -605,6 +614,7 @@ def main() -> None:
         sparse_anchor_indices,
         args.local_feature_root,
         args.local_modalities,
+        compute_transition_targets=args.transition_objective != "none",
     )
     local_contract = feature_metadata[0].get("local_contract")
     if not len(states):
@@ -613,11 +623,15 @@ def main() -> None:
         f"training examples: {len(states)}, feature_dim={hidden_size}, "
         f"input_layers={num_input_layers}"
     )
-    print(
-        "transition targets: "
-        f"{int(transition_targets.sum().item())}/{transition_targets.numel()} "
-        f"positive ({float(transition_targets.mean()) * 100:.3f}%)"
-    )
+    if args.transition_objective == "none":
+        print("transition targets: disabled (not derived or loaded)")
+    else:
+        print(
+            "transition targets: "
+            f"{int(transition_targets.sum().item())}/"
+            f"{transition_targets.numel()} positive "
+            f"({float(transition_targets.mean()) * 100:.3f}%)"
+        )
     if num_input_layers > 1 and args.layer_fusion == "single":
         raise ValueError(
             f"Features contain {num_input_layers} stacked layers; pass "
@@ -650,6 +664,7 @@ def main() -> None:
             args.local_feature_root,
             args.local_modalities,
             local_contract,
+            compute_transition_targets=args.transition_objective != "none",
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(