Setup the imports

kaseris · Feb 7, 2024 · 53b6b07 · 53b6b07
1 parent 5f34fa7
commit 53b6b07
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/skelcast/experiments/distributed.py b/src/skelcast/experiments/distributed.py
@@ -1,3 +1,7 @@
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel
+from torch.distributed import init_process_group, destroy_process_group
+
 from skelcast.experiments import RUNNERS
 
 @RUNNERS.register_module()