Have you tried the torch.split (https://pytorch.org/docs/stable/torch.html#torch.split) or torch.chunk (https://pytorch.org/docs/stable/torch.html#torch.chunk) function?
Have you tried the torch.split (https://pytorch.org/docs/stable/torch.html#torch.split) or torch.chunk (https://pytorch.org/docs/stable/torch.html#torch.chunk) function?