- and tuple(self.train_ar_mask.unique()) == (0, 1)
- and tuple(self.test_ar_mask.unique()) == (0, 1)
+ and tuple(x.item() for x in self.train_ar_mask.unique()) in { (0,), (1,), (0,1) }
+ and tuple(x.item() for x in self.test_ar_mask.unique()) in { (0,), (1,), (0,1) }