c = c.long()[:, None]
c = (
(1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
- * torch.tensor([128, 128, 128], device=c.device)
+ * torch.tensor([64, 64, 64], device=c.device)
+ (c == 1).long() * torch.tensor([0, 255, 0], device=c.device)
+ (c == 0).long() * torch.tensor([255, 255, 255], device=c.device)
+ (c == -1).long() * torch.tensor([255, 0, 0], device=c.device)