Function _drop_from_previous_rank

sat/vae_modules/cp_enc_dec.py:291–293

(input_, dim, kernel_size)

def _drop_from_previous_rank(input_, dim, kernel_size):
    input_ = input_.transpose(0, dim)[kernel_size - 1 :].transpose(0, dim)
    return input_


class _ConvolutionScatterToContextParallelRegion(torch.autograd.Function):