Error in training on cdsprites with ERM

[yongchanghao]

There is a filter of the indices that selects only latents[:, 0] != 3.

fish/src/models/datasets.py

Lines 266 to 271 in 333efa2

	if split=='val':
	self.latents = self.latents[domain_indices[-2]]
	self.images = self.images[domain_indices[-2]]
	elif split=='test':
	self.latents = self.latents[domain_indices[-1]]
	self.images = self.images[domain_indices[-1]]

In the code above, this condition is applied on val and test, but not on train. This results in an error in training with ERM. Could you provide the correct way to perform ERM on cdsprites?

Thanks.