Merge pull request #136 from GrzegorzKarchNV/master

Fixing concatenation error for fp16 distributed training
2019-02-01 12:10:42 -08:00 · 2019-02-01 12:10:42 -08:00 · f2c94d94fd
parent 825ffa47d1 df4a466af2
commit f2c94d94fd
1 changed files with 1 additions and 1 deletions
--- a/distributed.py
+++ b/distributed.py
@ -140,7 +140,7 @@ def apply_gradient_allreduce(module):
                buckets = {}
                for param in module.parameters():
                    if param.requires_grad and param.grad is not None:
-                        tp = type(param.data)
+                        tp = param.data.dtype
                        if tp not in buckets:
                            buckets[tp] = []
                        buckets[tp].append(param)