take out ocpdataparallel from fit.py

FAIR-Chem · Jan 4, 2024 · ddac40a · ddac40a
1 parent d3d7e1c
commit ddac40a
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/ocpmodels/modules/scaling/fit.py b/ocpmodels/modules/scaling/fit.py
@@ -10,7 +10,6 @@
 import torch.nn as nn
 from torch.nn.parallel.distributed import DistributedDataParallel
 
-from ocpmodels.common.data_parallel import OCPDataParallel
 from ocpmodels.common.flags import flags
 from ocpmodels.common.utils import (
     build_config,
@@ -78,7 +77,7 @@ def main(*, num_batches: int = 16) -> None:
         # unwrap module from DP/DDP
         unwrapped_model = model
         while isinstance(
-            unwrapped_model, (DistributedDataParallel, OCPDataParallel)
+            unwrapped_model, DistributedDataParallel
         ):
             unwrapped_model = unwrapped_model.module
         assert isinstance(