Use new TORCH_LIBRARY_FRAGMENT to register video_reder (#3105)

0a75a0c1 · Francisco Massa · GitHub · c5914452 · 0a75a0c1
Unverified Commit 0a75a0c1 authored Dec 03, 2020 by Francisco Massa Committed by GitHub Dec 03, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 9 deletions

torchvision/csrc/cpu/video_reader/VideoReader.cpp torchvision/csrc/cpu/video_reader/VideoReader.cpp +6 -9

No files found.
--- a/torchvision/csrc/cpu/video_reader/VideoReader.cpp
+++ b/torchvision/csrc/cpu/video_reader/VideoReader.cpp
@@ -661,12 +661,9 @@ torch::List<torch::Tensor> probeVideoFromFile(std::string videoPath) {

 } // namespace video_reader

-static auto registry = torch::RegisterOperators()
-                           .op("video_reader::read_video_from_memory",
-                               &video_reader::readVideoFromMemory)
-                           .op("video_reader::read_video_from_file",
-                               &video_reader::readVideoFromFile)
-                           .op("video_reader::probe_video_from_memory",
-                               &video_reader::probeVideoFromMemory)
-                           .op("video_reader::probe_video_from_file",
-                               &video_reader::probeVideoFromFile);
+TORCH_LIBRARY_FRAGMENT(video_reader, m) {
+  m.def("read_video_from_memory", video_reader::readVideoFromMemory);
+  m.def("read_video_from_file", video_reader::readVideoFromFile);
+  m.def("probe_video_from_memory", video_reader::probeVideoFromMemory);
+  m.def("probe_video_from_file", video_reader::probeVideoFromFile);
+}