pytorch - 💡(How to fix) Fix [torch.compile] Dynamo fails to trace assignment of custom tensor subclass to module attribute:self.custom_weights = CustomTensor(self.layer1.weight.data)`` [1 participants]

import torch
import torch.nn as nn
class CustomTensor(torch.Tensor):
    def __init__(self, data):
        super().__init__()
        self.custom_metadata = "custom_tensor_data"
    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        return super().__torch_function__(func, types, args, kwargs)
    def custom_method(self):
        return f"Custom tensor with metadata: {self.custom_metadata}"
class EnhancedCustomTensor(CustomTensor):
    def __init__(self, data):
        super().__init__(data)
        self.enhanced_flag = True
    def enhanced_operation(self):
        return self * 2.0 if self.enhanced_flag else self
class NeuralNetWithCustomTensors(nn.Module):
    def __init__(self, input_size=784, hidden_size=128, num_classes=10):
        super().__init__()
        self.input_size = input_size
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, hidden_size)
        self.output_layer = nn.Linear(hidden_size, num_classes)
        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
        self.custom_weights = None
        self.enhanced_features = None
    def forward(self, x):
        if not isinstance(x, CustomTensor):
            x = CustomTensor(x)
        self.custom_weights = CustomTensor(self.layer1.weight.data)
        x = self.layer1(x)
        x = self.activation(x)
        x = self.dropout(x)
        intermediate_result = CustomTensor(x)
        self.enhanced_features = EnhancedCustomTensor(intermediate_result)
        x = self.layer2(x)
        x = self.activation(x)
        if self.enhanced_features is not None:
            enhanced_result = self.enhanced_features.enhanced_operation()
            x = x + enhanced_result[:x.size(0), :x.size(1)]
        x = self.output_layer(x)
        return x
    def get_custom_metadata(self):
        if self.custom_weights is not None:
            return self.custom_weights.custom_method()
        return "No custom weights initialized"
def get_default_model():
    model = NeuralNetWithCustomTensors(input_size=784, hidden_size=128, num_classes=10)
    return model
def get_sample_inputs():
    batch_size = 4
    input_size = 784
    x = torch.randn(batch_size, input_size)
    return (x,)
def main():
    model = get_default_model()
    model.eval()
    inputs = get_sample_inputs()
    print("Model architecture:")
    print(model)
    print(f"\nCustom tensor subclass hierarchy: {model.get_custom_metadata()}")
    with torch.no_grad():
        output = model(*inputs)
    print("\nModel executed successfully!")
    print(f"Input shape: {inputs[0].shape}")
    print(f"Output shape: {output.shape}")
    print(f"Output type: {type(output)}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters())}")
    if model.enhanced_features is not None:
        print(f"\nEnhanced features type: {type(model.enhanced_features)}")
        print(f"Enhanced features metadata: {model.enhanced_features.custom_metadata}")
        print(f"Enhanced operation result shape: {model.enhanced_features.enhanced_operation().shape}")

    try:
        compiled_model = torch.compile(model)
        with torch.no_grad():
            output_compile = compiled_model(*inputs)
        print(f'Compile  shape: {output_compile.shape}')
    except Exception as e:
        print(f'\nOriginal exception:\n {e}')
if __name__ == "__main__":
    main()

---

Model architecture:
NeuralNetWithCustomTensors(
  (layer1): Linear(in_features=784, out_features=128, bias=True)
  (layer2): Linear(in_features=128, out_features=128, bias=True)
  (output_layer): Linear(in_features=128, out_features=10, bias=True)
  (activation): ReLU()
  (dropout): Dropout(p=0.2, inplace=False)
)

Custom tensor subclass hierarchy: No custom weights initialized

Model executed successfully!
Input shape: torch.Size([4, 784])
Output shape: torch.Size([4, 10])
Output type: <class '__main__.EnhancedCustomTensor'>
Model parameters: 118282

Enhanced features type: <class '__main__.EnhancedCustomTensor'>
Enhanced features metadata: custom_tensor_data
Enhanced operation result shape: torch.Size([4, 128])

Original exception:
 

from user code:
  ...ine 34, in torch_dynamo_resume_in_forward_at_33
    self.custom_weights = CustomTensor(self.layer1.weight.data)

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"