Is nn.Module ``stream-safe''?

autograd
1

Is it safe to share a module among different cuda streams during training?
More specifically, will module instance ``a’’ suffer from race condition in the following code?

form torch import nn
import torch

class A(nn.module):
    ...
class B(nn.module):
    ...
class C(nn.module):
    ...

class AB:
    def __init__(this,a,b):
        this.a=a;
        this.b=b;
        this.stream=torch.cuda.Stream();
    def involke(x):
        with torch.cuda.stream(this.stream):
            r=this.b(this.a(x));
    def wait():
        this.stream.synchronize();


class AC:
    def __init__(this,a,c):
        this.a=a;
        this.c=c;
        this.stream=torch.cuda.Stream();
    def involke(x):
        with torch.cuda.stream(this.stream):
            r=this.c(this.a(x));
    def wait():
        this.stream.synchronize();


def main():
    a,b,c=A(),B(),C()
    fab=AB(a,b);
    fac=AC(a,c);
    ...
    x=loaddata();
    y_ab=fab.involke(x);
    y_ac=fac.involke(x);
    fab.wait();
    fac.wait();
    ...