diff --git a/sleap/nn/tracking.py b/sleap/nn/tracking.py
index 746ab2881..d79c67d75 100644
--- a/sleap/nn/tracking.py
+++ b/sleap/nn/tracking.py
@@ -667,6 +667,7 @@ def run_tracker(
     def track(
         self,
         untracked_instances: List[InstanceType],
+        img_hw: Tuple[int],
         img: Optional[np.ndarray] = None,
         t: int = None,
     ):
@@ -1561,6 +1562,7 @@ def cull_function(inst_list):
     def track(
         self,
         untracked_instances: List[InstanceType],
+        img_hw: Tuple[int],
         img: Optional[np.ndarray] = None,
         t: int = None,
     ) -> List[InstanceType]:
diff --git a/tests/nn/test_tracker_components.py b/tests/nn/test_tracker_components.py
index 9d3b65b38..fa0cc5f51 100644
--- a/tests/nn/test_tracker_components.py
+++ b/tests/nn/test_tracker_components.py
@@ -9,7 +9,6 @@
     FrameMatches,
     greedy_matching,
 )
-from sleap.io.dataset import Labels
 
 from sleap.instance import PredictedInstance
 from sleap.skeleton import Skeleton
diff --git a/tests/nn/test_tracking_integration.py b/tests/nn/test_tracking_integration.py
index caebe49ff..c479462f8 100644
--- a/tests/nn/test_tracking_integration.py
+++ b/tests/nn/test_tracking_integration.py
@@ -1,4 +1,3 @@
-import inspect
 import operator
 import os
 import time
@@ -7,7 +6,7 @@
 import sleap
 from sleap.nn.inference import main as inference_cli
 import sleap.nn.tracker.components
-from sleap.io.dataset import Labels, LabeledFrame
+from sleap.io.dataset import Labels
 
 
 def test_simple_tracker(tmpdir, centered_pair_predictions_slp_path):