Ryu1845

import torch
import numpy as np

# numpy
a = np.random.rand(10, 20)
tmp0 = np.split(a, indices_or_sections=5, axis=0) # split into 5 sections
for t in tmp0:
    print(t.shape)
# (2, 20)
# (2, 20)