Pass Stream as arg to torch.jit.script function

jit
1

I can’t pass cuda stream as an arg to torch.jit.script function.
Following is code snippet.
Any help?

  @torch.jit.script
  def f(d, s: torch.Stream, i: int):
    with torch.cuda.stream(s):
        #do sth

  @torch.jit.script
  def r(tensors: List[Tensor], num_gpus: int):
    futures: List[torch.jit.Future[Tensor]] = []
    streams = [
        torch.cuda.Stream(torch.device(f"cuda:{i}"))
        for i in range(num_gpus)
    ]
    for i in range(num_gpus):
      futures.append(torch.jit.fork(f, tensors[i], streams[i], i))
    results = [torch.jit.wait(fut) for fut in futures]
    return results

Got RuntimeError:

f(Tensor d, Stream s, int i) -> (Tensor):
Expected a value of type 'Stream' for argument 's' but instead found type '__torch__.torch.classes.cuda.Stream (of Python compilation unit at: 0)'.

I also tried typing torch.cuda.Stream, and got:

'__torch__.torch.cuda.streams.Stream (of Python compilation unit at: 0x3e98e70)' object has no attribute or method 'cuda_stream'. Did you forget to initialize an attribute in __init__()?:
  File "/data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/cuda/streams.py", line 98
    @property
    def _as_parameter_(self):
        return ctypes.c_void_p(self.cuda_stream)
                               ~~~~~~~~~~~~~~~~ <--- HERE
'Stream.___as_parameter__getter' is being compiled since it was called from '__torch__.torch.cuda.streams.Stream'