wrap_fx_proxy_cls doesn't cover torch.cuda.Stream about torchdynamo HOT 3 CLOSED

import torch
import torch.nn as nn
import torch._dynamo as dynamo
from torch.nn.parallel._functions import _get_stream

class CopyModel(nn.Module):
    def __init__(self):
        super(CopyModel, self).__init__()

    def forward(self, input, device):
        self.copy_stream = torch.cuda.Stream(device)
        with torch.cuda.stream(self.copy_stream):
            return input.cuda()

model = CopyModel()
output = model(torch.rand(1000), 1)
print(output.get_device())

opt_model = dynamo.optimize('eager')(model)
output = opt_model(torch.rand(1000), 1)
print(output.get_device())