import torch._dynamo as dynamo
import torch

dynamo.config.log_level = dynamo.config.logging.DEBUG
torch._dynamo.reset()

def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
    print(gm.code)
    print(gm.graph)
    gm.graph.print_tabular()
    return gm.forward

s = torch.cuda.Stream()

@dynamo.optimize(my_compiler)
def fn(t) -> torch.Tensor:
    tmp1 = torch.mul(t, 5)
    tmp2 = torch.add(tmp1, 2)
    with torch.cuda.stream(s):
      r = torch.relu(tmp2)
    return r

i = torch.Tensor([-2, 3]).to('cuda')
r = fn(i)
print(f"r = {r}")

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/utils.py", line 420, in proxy_args_kwargs
    proxy_args = tuple(arg.as_proxy() for arg in args)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/utils.py", line 420, in <genexpr>
    proxy_args = tuple(arg.as_proxy() for arg in args)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/base.py", line 206, in as_proxy
    raise NotImplementedError(str(self))
NotImplementedError: UserDefinedObjectVariable(Stream)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 307, in wrapper
    return inner_fn(self, inst)
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 974, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/symbolic_convert.py", line 435, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/variables/torch.py", line 476, in call_function
    *proxy_args_kwargs(args, kwargs),
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/utils.py", line 427, in proxy_args_kwargs
    raise unimplemented(
  File "/usr/local/lib/python3.8/dist-packages/torch/_dynamo/exc.py", line 71, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: call_function args: UserDefinedObjectVariable(Stream) 
[2023-01-23 04:54:12,810] torch._dynamo.output_graph: [DEBUG] restore_graphstate: removed 0 nodes
[2023-01-23 04:54:12,811] torch._dynamo.output_graph: [DEBUG] COMPILING GRAPH due to GraphCompileReason(reason='call_function args: UserDefinedObjectVariable(Stream) ', user_stack=[<FrameSummary file <ipython-input-12-f1436f0fcd91>, line 20 in fn>])

def forward(self, t : torch.Tensor):
    mul = torch.mul(t, 5);  t = None
    add = torch.add(mul, 2);  mul = None
    return (add,)
    
graph():
    %t : torch.Tensor [#users=1] = placeholder[target=t]
    %mul : [#users=1] = call_function[target=torch.mul](args = (%t, 5), kwargs = {})
    %add : [#users=1] = call_function[target=torch.add](args = (%mul, 2), kwargs = {})
    return (add,)
opcode         name    target                                                  args       kwargs
-------------  ------  ------------------------------------------------------  ---------  --------
placeholder    t       t                                                       ()         {}
call_function  mul     <built-in method mul of type object at 0x7f2e333dbe40>  (t, 5)     {}
call_function  add     <built-in method add of type object at 0x7f2e333dbe40>  (mul, 2)   {}
output         output  output                                                  ((add,),)  {}