Merge pull request #32 from lucasnewman/fix-import

Fix DDP import
lucidrains · Oct 6, 2023 · d8c59f1 · d8c59f1
2 parents 7ce0e56 + f25f713
commit d8c59f1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/voicebox_pytorch/trainer.py b/voicebox_pytorch/trainer.py
@@ -14,6 +14,7 @@
 from voicebox_pytorch.optimizer import get_optimizer
 
 from accelerate import Accelerator, DistributedType
+from accelerate.utils import DistributedDataParallelKwargs
 
 # helpers