How make customised dataset for semantic segmentation?

Yes, transforms.ToTensor will give you an image tensor with values in the range [0, 1].
This might be sufficient to train your model, however usually you would standardize your tensors to have zero-mean and a stddev of 1. You could calculate the mean and stddev of your train images yourself using this small example or alternatively the ImageNet mean and std work quite well for normal images.

The code looks good! You woun’t need the train argument anymore in your Dataset, but besides that I cannot find anything obviously wrong.

3 Likes