Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "a799ca343b13665661a5e95f5ad1523457bef2e2"
Unverified Commit 72bd7c69 authored by Frank Lee's avatar Frank Lee Committed by GitHub
Browse files

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

parent 5a9d8ef4
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
colossalai/amp/naive_amp/naive_amp.py
View file @ 72bd7c69
......@@ -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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment