handle out != None

danielvegamyhre · danielvegamyhre · commit 3201f9ef509b · 2025-07-02T09:48:57.000-07:00
diff --git a/torchao/prototype/moe_training/tensor.py b/torchao/prototype/moe_training/tensor.py
@@ -183,6 +183,7 @@ def fsdp_post_all_gather(
                     f"{out_data.dtype} {param_dtype}"
                 )
                 out_data.copy_(data)
+
             return
 
         # For training step 0, out=None, so we need to return a new ScaledGroupedMMTensor.

Original file line number	Diff line number	Diff line change
`@@ -183,6 +183,7 @@ def fsdp_post_all_gather(`
`183`	`183`	`f"{out_data.dtype} {param_dtype}"`
`184`	`184`	`)`
`185`	`185`	`out_data.copy_(data)`
	`186`	`+`
`186`	`187`	`return`
`187`	`188`
`188`	`189`	`# For training step 0, out=None, so we need to return a new ScaledGroupedMMTensor.`