4bcf6fdd93291a1bfba3c7038fe30a7338870968,onmt/models/SRU.py,,,#,72

Before Change

    from pynvrtc.compiler import Program

    // This sets up device to use.
    device = torch.device("cuda")
    tmp_ = torch.rand(1, 1).to(device)

    sru_prog = Program(SRU_CODE.encode("utf-8"),
                       "sru_prog.cu".encode("utf-8"))
    sru_ptx = sru_prog.compile()
    sru_mod = function.Module()
    sru_mod.load(bytes(sru_ptx.encode()))

    SRU_FWD_FUNC = sru_mod.get_function("sru_fwd")
    SRU_BWD_FUNC = sru_mod.get_function("sru_bwd")
    SRU_BiFWD_FUNC = sru_mod.get_function("sru_bi_fwd")
    SRU_BiBWD_FUNC = sru_mod.get_function("sru_bi_bwd")

    stream = namedtuple("Stream", ["ptr"])
    SRU_STREAM = stream(ptr=torch.cuda.current_stream().cuda_stream)


class SRU_Compute(Function):

After Change

}

SRU_FWD_FUNC, SRU_BWD_FUNC = None, None
SRU_BiFWD_FUNC, SRU_BiBWD_FUNC = None, None
SRU_STREAM = None

Instances