github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-17 07:26:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Frank Lee 2022-06-13 14:18:04 +08:00 committed by GitHub
parent 5a9d8ef4d5
commit 72bd7c696b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/amp/naive_amp/naive_amp.py
View File

@ -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