assert outputs.is_cuda or outputs.device.type == 'xla'

[morestart]

    def training_step(self, batch, batch_idx):
        x, y = batch['input'], batch['target']
        predict_air, predict_surface = self.forward(x[0], x[1])
        loss_dict = lat_weighted_mae(predict_air, predict_surface, y[0], y[1], self.lat,
                                     self.mean, self.std, self.upper_air_vars,
                                     self.surface_vars)
        for var in loss_dict.keys():
            self.log(
                "train/" + var,
                loss_dict[var],
                on_step=True,
                on_epoch=False,
                prog_bar=True,
            )
        return loss_dict['w_mae']

    def validation_step(self, batch, batch_idx):
        x, y = batch['input'], batch['target']
        predict_air, predict_surface = self.forward(x[0], x[1])
        loss_dict = lat_weighted_mae_val(predict_air, predict_surface, y[0], y[1], self.lat,
                                         self.mean, self.std, self.upper_air_vars,
                                         self.surface_vars)
        for var in loss_dict.keys():
            self.log(
                "val/" + var,
                loss_dict[var],
                on_step=False,
                on_epoch=True,
                prog_bar=False,
                sync_dist=True,
            )
        return loss_dict

w_mae: tensor(78.9927, grad_fn=)