'DataLoader' is not defined

[Hnd7]

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]

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