def METHOD_NAME(self, *args, **kwargs):
    saved_inference = self.inference
    self.inference = True
    dataset = super().METHOD_NAME(*args, **kwargs)
    self.inference = saved_inference
    return dataset