(self, input_shape, num_classes, num_domains, hparams)
| 307 | """Invariant Risk Minimization""" |
| 308 | |
| 309 | def __init__(self, input_shape, num_classes, num_domains, hparams): |
| 310 | super(IRM, self).__init__(input_shape, num_classes, num_domains, hparams) |
| 311 | self.register_buffer('update_count', torch.tensor([0])) |
| 312 | |
| 313 | @staticmethod |
| 314 | def _irm_penalty(logits, y): |