preserve subclass with transpose

danielvegamyhre · danielvegamyhre · commit 15f381e15062 · 2025-08-01T14:08:13.000-07:00
diff --git a/torchao/prototype/moe_training/tensor.py b/torchao/prototype/moe_training/tensor.py
@@ -30,6 +30,7 @@
     torch.ops.aten._pin_memory.default,
     torch.ops.aten.split.Tensor,
     torch.ops.aten.clone.default,
+    torch.ops.aten.transpose.int,
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@`
`30`	`30`	`torch.ops.aten._pin_memory.default,`
`31`	`31`	`torch.ops.aten.split.Tensor,`
`32`	`32`	`torch.ops.aten.clone.default,`
	`33`	`+ torch.ops.aten.transpose.int,`
`33`	`34`	`}`
`34`	`35`
`35`	`36`