Method forward

sat/sgm/modules/cp_enc_dec.py:334–340

(self, input_)

class ContextParallelGroupNorm(torch.nn.GroupNorm):
    def forward(self, input_):
        if _USE_CP:
            input_ = conv_gather_from_context_parallel_region(input_, dim=2, kernel_size=1)
        output = super().forward(input_)
        if _USE_CP:
            output = conv_scatter_to_context_parallel_region(output, dim=2, kernel_size=1)
        return output


def Normalize(in_channels, gather=False, **kwargs):  # same for 3D and 2D