Fix fsd_model/cot_reasoning.py for training (autograd + in-place ops)

fsd_model/cot_reasoning.py

@@ -491,21 +491,25 @@ class SafetyDecisionGate(nn.Module):
         planner_speeds = planner_waypoints[:, :, 3]
         cot_speeds = cot_waypoints[:, :, 3]
         safe_speeds = torch.min(planner_speeds, F.relu(cot_speeds))
-        cot_waypoints = cot_waypoints.clone()
-        cot_waypoints[:, :, 3] = safe_speeds
-
-        # Clamp all speeds
-        cot_waypoints[:, :, 3] = torch.clamp(cot_waypoints[:, :, 3], 0.0, self.max_speed_ms)
+        safe_speeds = torch.clamp(safe_speeds, 0.0, self.max_speed_ms)
+        
+        # Build cot_waypoints without in-place ops
+        cot_waypoints = torch.cat([
+            cot_waypoints[:, :, :3],
+            safe_speeds.unsqueeze(-1),
+        ], dim=-1)
 
         # Blend: output = (1-alpha)*planner + alpha*cot
         alpha_expanded = alpha.unsqueeze(-1)  # (B, 1, 1)
         gated_waypoints = (1 - alpha_expanded) * planner_waypoints + alpha_expanded * cot_waypoints
 
         # Ensure gated speeds never exceed planner speeds (monotonic safety)
-        gated_waypoints[:, :, 3] = torch.min(
-            gated_waypoints[:, :, 3], planner_waypoints[:, :, 3]
-        )
-        gated_waypoints[:, :, 3] = torch.clamp(gated_waypoints[:, :, 3], 0.0, self.max_speed_ms)
+        gated_speeds = torch.min(gated_waypoints[:, :, 3], planner_waypoints[:, :, 3])
+        gated_speeds = torch.clamp(gated_speeds, 0.0, self.max_speed_ms)
+        gated_waypoints = torch.cat([
+            gated_waypoints[:, :, :3],
+            gated_speeds.unsqueeze(-1),
+        ], dim=-1)
 
         # Post-gate safety score
         safety = self.safety_score(justification_embedding)