# check shapes 
images, masks = next(iter(train_loader))
print(images.shape)
print(masks.shape)

# %%
print(len(train_loader))

# %%
images, masks = next(iter(test_loader))
print(images.shape)
print(masks.shape)

# %%
print((len(test_loader)))

#[batch_size, no_of_classes, size, size]

torch.Size([8, 3, 320, 320])
torch.Size([8, 12, 320, 320])
46
torch.Size([8, 3, 384, 480])
torch.Size([8, 12, 384, 480])
13