import torch
from torch.utils.data import TensorDataset, DataLoader
X=torch.arange(22,dtype=torch.float32).reshape(11,2)
y=torch.arange(11,dtype=torch.float32).reshape(11,1)
loader=DataLoader(TensorDataset(X,y),batch_size=4,shuffle=False,num_workers=0)
sizes=[]; total=0.; count=0
for xb,yb in loader:
    assert xb.shape[1]==2 and yb.shape==(len(xb),1)
    sizes.append(len(xb))
    mean=yb.square().mean().item()
    total+=mean*len(xb); count+=len(xb)
assert sizes==[4,4,3] and count==11
assert abs(total/count-y.square().mean().item())<1e-6
print('batch_sizes=',sizes,'sample_count=',count,'mean_target_square=',total/count)