search

mberkay0 / pretrained-backbones-unet

A PyTorch-based Python library with UNet architecture and multiple backbones for Image Semantic Segmentation.
MIT License
46 stars 9 forks source link

'DataLoader' is not defined #6

Open Hnd7 opened 9 months ago

Hnd7 commented 9 months ago

Hi, when I am trying to start the training I am getting the following error:

NameError                                 Traceback (most recent call last)
/tmp/ipykernel_21140/3294596851.py in <module>
     14 val_dataset = SemanticSegmentationDataset(val_img_path, val_mask_path)
     15 
---> 16 train_loader = DataLoader(train_dataset, batch_size=2)
     17 val_loader = DataLoader(val_dataset, batch_size=2)
     18 

NameError: name 'DataLoader' is not defined

Do you have any suggestions to how to solve it?

Thanks!

AlexLee1235 commented 7 months ago

I think you can use from torch.utils.data import Dataset, DataLoader