- assert (
- self.nb_codes <= max_nb_codes
- and self.train_input.min() >= 0
- and self.test_input.min() >= 0
- and tuple(self.train_ar_mask.unique()) == (0, 1)
- and tuple(self.test_ar_mask.unique()) == (0, 1)
- )
+ assert self.nb_codes <= max_nb_codes
+ assert self.train_input.min() >= 0
+ assert self.test_input.min() >= 0
+ assert tuple(x.item() for x in self.train_ar_mask.unique()) in {
+ (0,),
+ (1,),
+ (0, 1),
+ }
+ assert tuple(x.item() for x in self.test_ar_mask.unique()) in {
+ (0,),
+ (1,),
+ (0, 1),
+ }