Fix some stuff?

Author: NicolasHug

test/decoders/test_video_decoder.py

@@ -15,8 +15,8 @@
     assert_frames_equal,
     AV1_VIDEO,
     cpu_and_cuda,
+    get_ffmpeg_major_version,
     H265_VIDEO,
-    in_fbcode,
     NASA_VIDEO,
 )
 
@@ -249,7 +249,7 @@ def test_getitem_slice(self, device, seek_mode):
             ]
         )
         for sliced, ref in zip(all_frames, decoder):
-            if not (in_fbcode() and device == "cuda"):
+            if not (device == "cuda" and get_ffmpeg_major_version() == 4):
                 # TODO: remove the "if".
                 # See https://github.com/pytorch/torchcodec/issues/428
                 assert_frames_equal(sliced, ref)
@@ -428,7 +428,7 @@ def test_get_frames_at_fails(self, device, seek_mode):
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_get_frame_at_av1(self, device):
-        if in_fbcode() and device == "cuda":
+        if device == "cuda" and get_ffmpeg_major_version() == 4:
             return
 
         decoder = VideoDecoder(AV1_VIDEO.path, device=device)

test/utils.py

@@ -25,6 +25,10 @@ def cpu_and_cuda():
     return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
 
 
+def get_ffmpeg_major_version():
+    return int(get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0])
+
+
 # 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
@@ -33,11 +37,8 @@ def cpu_and_cuda():
 def assert_frames_equal(*args, **kwargs):
     if sys.platform == "linux":
         if args[0].device.type == "cuda":
-            ffmpeg_major_version = int(
-                get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0]
-            )
             atol = 2
-            if in_fbcode() or ffmpeg_major_version == 4:
+            if get_ffmpeg_major_version() == 4:
                 assert_tensor_close_on_at_least(
                     args[0], args[1], percentage=95, atol=atol
                 )