bug for color_seg

[zhuyb00]

When I run the following pretrain code python train_segmentation.py attention_mr

It turns out the error that

line 37, in color_seg out.append((colors[s[0]]).permute(2, 0, 1)) RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

then I change the code to

def color_seg(seg): out = [] seg = seg.cpu for s in seg: out.append((colors[s[0]]).permute(2, 0, 1)) return torch.stack(out, dim=0)

Rerun the code turns out that

line 38, in color_seg out.append((colors[s[0]]).permute(2, 0, 1)) TypeError: 'builtin_function_or_method' object is not subscriptable