fp16 fix

NoahSchiro · NoahSchiro · commit 7bac936a5a95 · 2024-03-11T18:45:34.000-04:00
diff --git a/improved_diffusion/fp16_util.py b/improved_diffusion/fp16_util.py
@@ -65,7 +65,7 @@ def unflatten_master_params(model_params, master_params):
     """
     Unflatten the master parameters to look like model_params.
     """
-    return _unflatten_dense_tensors(master_params[0].detach(), model_params)
+    return _unflatten_dense_tensors(master_params[0].detach(), tuple(tensor for tensor in model_params))
 
 
 def zero_grad(model_params):
