Load and importlib

Author: scotts

src/torchcodec/decoders/_core/CMakeLists.txt

@@ -97,11 +97,6 @@ function(make_torchcodec_libraries
         "${decoder_library_name}"
         "${ffmpeg_include_dirs}"
     )
-    target_compile_definitions(
-        ${pybind_ops_library_name}
-        PUBLIC
-        TORCHCODEC_PYBIND=_torchcodec_pybind_ops${ffmpeg_major_version}
-    )
     # pybind11 quirk, see:
     #   https://pybind11.readthedocs.io/en/stable/faq.html#someclass-declared-with-greater-visibility-than-the-type-of-its-field-someclass-member-wattributes
     target_compile_options(

src/torchcodec/decoders/_core/PyBindOps.cpp

@@ -37,11 +37,7 @@ int64_t create_from_file_like(
   return reinterpret_cast<int64_t>(decoder);
 }
 
-#ifndef TORCHCODEC_PYBIND
-#error TORCHCODEC_PYBIND must be defined.
-#endif
-
-PYBIND11_MODULE(TORCHCODEC_PYBIND, m) {
+PYBIND11_MODULE(_torchcodec_pybind_ops, m) {
   m.def("create_from_file_like", &create_from_file_like);
 }
 

src/torchcodec/decoders/_core/ops.py

@@ -35,17 +35,18 @@ def load_torchcodec_extension():
         decoder_library_name = f"libtorchcodec_decoder{ffmpeg_major_version}"
         custom_ops_library_name = f"libtorchcodec_custom_ops{ffmpeg_major_version}"
         pybind_ops_library_name = f"libtorchcodec_pybind_ops{ffmpeg_major_version}"
-        pybind_ops_module_name = f"_torchcodec_pybind_ops{ffmpeg_major_version}"
+        pybind_ops_module_name = "torchcodec._torchcodec_pybind_ops"
         try:
             torch.ops.load_library(_get_extension_path(decoder_library_name))
             torch.ops.load_library(_get_extension_path(custom_ops_library_name))
+            torch.ops.load_library(_get_extension_path(pybind_ops_library_name))
 
             spec = importlib.util.spec_from_file_location(
                 pybind_ops_module_name,
                 _get_extension_path(pybind_ops_library_name),
             )
             if spec is None:
-                raise ImportError(f"Unable to load spec for pybind_ops {_get_extension_path(pybind_ops_library_name)}")
+                raise ImportError(f"Unable to load spec for pybind_ops")
 
             global _pybind_ops
             _pybind_ops = importlib.util.module_from_spec(spec)