fix FP16 bug of track_id annotation

ifzhang · Apr 27, 2022 · 7d71238 · 7d71238
1 parent f137cbf
commit 7d71238
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/yolox/core/trainer.py b/yolox/core/trainer.py
@@ -91,6 +91,8 @@ def train_one_iter(self):
         iter_start_time = time.time()
 
         inps, targets = self.prefetcher.next()
+        track_ids = targets[:, :, 5]
+        targets = targets[:, :, :5]
         inps = inps.to(self.data_type)
         targets = targets.to(self.data_type)
         targets.requires_grad = False