Fix some tests

Author: NicolasHug

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -1017,7 +1017,10 @@ void VideoDecoder::maybeSeekToBeforeDesiredPts() {
     //   exactly at a frame start either.
     // - When calling `getFramePlayedAt(pts)`, regardless of the mode, if `pts`
     //   doesn't land exactly at a frame's start. We have tests that currently
-    //   exhibit this behavior: test_get_frame_at_pts_audio_bad().
+    //   exhibit this behavior: test_get_frame_at_pts_audio_bad(). The "obvious"
+    //   fix for this is to let `getFramePlayedAt` convert the pts to an index,
+    //   just like the rest of the APIs.
+    //
     // TODO HOW DO WE FIX THIS??
 
     // A few notes:

test/utils.py

@@ -23,12 +23,25 @@ def cpu_and_cuda():
     return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
 
 
+def assert_frames_equal(*args, **kwargs):
+    frame = args[0]
+    # This heuristic will work until we start returningu int8 audio frames...
+    if frame.dtype == torch.uint8:
+        return assert_video_frames_equal(*args, **kwargs)
+    else:
+        return assert_audio_frames_equal(*args, **kwargs)
+
+
+def assert_audio_frames_equal(*args, **kwargs):
+    torch.testing.assert_close(*args, **kwargs)
+
+
 # For use with decoded data frames. On CPU Linux, we expect exact, bit-for-bit
 # equality. On CUDA Linux, we expect a small tolerance.
 # On other platforms (e.g. MacOS), we also allow a small tolerance. FFmpeg does
 # not guarantee bit-for-bit equality across systems and architectures, so we
 # also cannot. We currently use Linux on x86_64 as our reference system.
-def assert_frames_equal(*args, **kwargs):
+def assert_video_frames_equal(*args, **kwargs):
     if sys.platform == "linux":
         if args[0].device.type == "cuda":
             atol = 2