(self, input_shape, num_classes, num_domains, hparams)
| 176 | """ Adaptive Risk Minimization (ARM) """ |
| 177 | |
| 178 | def __init__(self, input_shape, num_classes, num_domains, hparams): |
| 179 | original_input_shape = input_shape |
| 180 | input_shape = (1 + original_input_shape[0],) + original_input_shape[1:] |
| 181 | super(ARM, self).__init__(input_shape, num_classes, num_domains, hparams) |
| 182 | self.context_net = networks.ContextNet(original_input_shape) |
| 183 | self.support_size = hparams['batch_size'] |
| 184 | |
| 185 | def predict(self, x): |
| 186 | batch_size, c, h, w = x.shape |