So I made my dataset the way you described it, I put num_classes to 3 and changed the labels line.
I now get an error in the following part of the code:
# split the color-encoded mask into a set of binary masks

  masks = mask == obj_ids[:, None, None]
  print(obj_ids[:, None, None])
        # get bounding box coordinates for each mask
        num_objs = len(obj_ids)

        boxes = []
        for i in range(num_objs):
            pos = np.where(masks[i])
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
            boxes.append([xmin, ymin, xmax, ymax])

The error that I get:
    data.reraise()
  File "anaconda3/lib/python3.7/site-packages/torch/_utils.py", line 394, in reraise
    raise self.exc_type(msg)
TypeError: Caught TypeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "anaconda3/lib/python3.7/site-packages/torch/utils/data/_utils/worker.py", line 178, in _worker_loop
    data = fetcher.fetch(index)
  File "anaconda3/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "anaconda3/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "multi_training.py", line 55, in __getitem__
    pos = np.where(masks[i])
TypeError: 'bool' object is not subscriptable

The obj_ids print statement gives:
[[[1]]

 [[2]]]

masks become False instead of a 2d list with false and true
What can I do? or did I made a mistake with the masks?
This is one of the masks: (note, iris = (1,1,1) pupil = (2,2,2))

The changes are working with the old, 2 class dataset.