- intact = seq[:, : height * width].reshape(-1, height, width)
- masked = seq[:, height * width + 1 :].reshape(-1, height, width)
- img_intact, img_masked = colors[intact], colors[masked]
+ f_start = seq[:, : height * width].reshape(-1, height, width)
+ f_start = (f_start >= len(colors)).long() + (f_start < len(colors)).long() * f_start
+ f_end = seq[:, height * width + 1 :].reshape(-1, height, width)
+ f_end = (f_end >= len(colors)).long() + (f_end < len(colors)).long() * f_end
+
+ img_f_start, img_f_end = colors[f_start], colors[f_end]