Fix save error for multi-gpu

2020-07-04 16:50:33 +07:00 · 2020-07-04 16:50:33 +07:00 · f83521ab9a
parent 5e2429e618
commit f83521ab9a
1 changed files with 1 additions and 1 deletions
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@ -201,5 +201,5 @@ class ModelEMA:
    def update_attr(self, model):
        # Update EMA attributes
        for k, v in model.__dict__.items():
-            if not k.startswith('_') and k != 'module':
+            if not k.startswith('_') and k not in ["module", "process_group", "reducer"]:
                setattr(self.ema, k, v)