[amp] included dict for type casting of model output (#1102)

2022-06-13 14:18:04 +08:00 · 2022-06-13 14:18:04 +08:00 · 72bd7c696b
parent 5a9d8ef4d5
commit 72bd7c696b
1 changed files with 2 additions and 0 deletions
--- a/colossalai/amp/naive_amp/naive_amp.py
+++ b/colossalai/amp/naive_amp/naive_amp.py
@ -149,4 +149,6 @@ class NaiveAMPModel(nn.Module):
                out = self._convert_to_fp32(out)
            elif isinstance(out, (tuple, list)):
                out = [self._convert_to_fp32(val) for val in out]
+            elif isinstance(out, dict):
+                out = {key: self._convert_to_fp32(val) for key, val in out.items()}
        return out