Fix ddp bug

2021-06-13 10:25:23 -06:00 · 2021-06-13 10:25:23 -06:00 · 1cd75dfd33
commit 1cd75dfd33
parent 3e3ad7825f
1 changed files with 3 additions and 2 deletions
--- a/codes/trainer/base_model.py
+++ b/codes/trainer/base_model.py
@ -98,7 +98,8 @@ class BaseModel():
        return save_path

    def load_network(self, load_path, network, strict=True, pretrain_base_path=None):
-        #if isinstance(network, nn.DataParallel) or isinstance(network, DistributedDataParallel):
+        # Sometimes networks are passed in as DDP modules, we want the raw parameters.
+        if hasattr(network, 'module'):
            network = network.module
        load_net = torch.load(load_path)