1 | #pragma once |
---|---|
2 | |
3 | #include <string> |
4 | #include <vector> |
5 | #include <cassert> |
6 | #include <memory> |
7 | |
8 | #include <ATen/DynamicLibrary.h> |
9 | |
10 | namespace c10d { |
11 | |
12 | inline std::shared_ptr<at::DynamicLibrary> loadTorchUCC() { |
13 | const char *path = std::getenv("TORCH_UCC_LIBRARY_PATH"); |
14 | if (path != nullptr) { |
15 | try { |
16 | return std::make_shared<at::DynamicLibrary>(path); |
17 | } catch (const c10::DynamicLibraryError &e) { |
18 | TORCH_WARN("TORCH_UCC_LIBRARY_PATH is set, " |
19 | "but the loading of torch_ucc.so failed with:", e.msg()); |
20 | } |
21 | } |
22 | return nullptr; |
23 | } |
24 | |
25 | } // namespace c10d |
26 |