code: def test_await_class_return(self):
      class C:
          __slots__ = ["a", "b"]
  
          def __init__(self, a: Tensor, b: Tensor):
              self.a = a
              self.b = b
  
      make_global(C)
  
      # Can not stay in the class as Jit does not support Recursive annotations
      # (self in wait_impl can not be annotated as C as C is not defined by this time)
      def C_wait_impl(self: C) -> C:
          return C(self.a * 2, self.b * 3)
  
      def fn_arg_C(x: C) -> Tensor:
          return x.a + x.b
  
      def fn(x: Tensor):
          aw: Await[C] = torch.jit._awaitable(C_wait_impl, C(x, x))
          _a = torch.eye(2)
          y = fn_arg_C(torch.jit._awaitable_wait(aw))
          return _a + y + x
  
      inp = torch.ones(2)
  
      sm = torch.jit.script(fn)
      out = fn(inp)
      script_out = sm(inp)
      self.assertTrue(
          torch.allclose(torch.eye(2) + 6 * torch.ones(2), script_out)
      )
      self.assertTrue(torch.allclose(script_out, out))
      self.assertGraphContainsExactly(
          sm.graph, kind="prim::awaitable_wait", num_kind_nodes=1
      )
  
imports: import io
  import torch
  from torch.testing._internal.jit_utils import JitTestCase
  from torch.testing._internal.jit_utils import make_global
  from typing import List, Optional, Tuple
  from torch import Tensor
  from torch._awaits import _Await
label_desc: Use image node_id 8 for calling the TestAwait obj's underlying member method code with example usage: obj.test_await_class_return() and returns: C, unknown, unknown