# Buggy code
def _run_early_stopping_check(self, trainer, pl_module):
    logs = trainer.callback_metrics
    if not self._validate_condition_metric(logs):
        return  # short circuit if metric not present

    current = logs.get(self.monitor)
    if not isinstance(current, torch.Tensor):
        current = torch.tensor(current)

    if self.monitor_op(current - self.min_delta, self.best_score):
        self.best_score = current
        self.wait_count = 0
    else:
        self.wait_count += 1
        if self.wait_count >= self.patience:
            self.stopped_epoch = trainer.current_epoch
            trainer.should_stop = True


# Traceback (RuntimeError)
Exception in device=TPU:2: torch_xla/csrc/aten_xla_bridge.cpp:69 : Check failed: xtensor
*** Begin stack trace ***
tensorflow::CurrentStackTrace[abi:cxx11]()
torch_xla::bridge::GetXlaTensor(at::Tensor const&amp;)
torch_xla::AtenXlaType::lt(at::Tensor const&amp;, at::Tensor const&amp;)
c10::detail::wrap_kernel_functor_unboxed_<c10::detail::WrapRuntimeKernelFunctor_<at::Tensor (*)(at::Tensor const&amp;, at::Tensor const&amp;), at::Tensor, c10::guts::typelist::typelist<at::Tensor const&amp;, at::Tensor const&amp;> >, at::Tensor (at::Tensor const&amp;, at::Tensor const&amp;)>::call(c10::OperatorKernel*, at::Tensor const&amp;, at::Tensor const&amp;)

_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyFunction_FastCallDict

PyObject_Call
_PyEval_EvalFrameDefault

PyObject_Call
_PyEval_EvalFrameDefault

PyObject_Call
_PyEval_EvalFrameDefault

_PyEval_EvalFrameDefault

_PyEval_EvalFrameDefault

_PyEval_EvalFrameDefault
_PyFunction_FastCallDict

_PyObject_FastCallKeywords

_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault

_PyFunction_FastCallDict

PyObject_Call
_PyEval_EvalFrameDefault

_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault

PyObject_Call
_PyEval_EvalFrameDefault

PyObject_Call
_PyEval_EvalFrameDefault

_PyEval_EvalFrameDefault

*** End stack trace ***
Input tensor is not an XLA tensor: torch.FloatTensor
Traceback (most recent call last):
File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py", line 119, in _start_fn
fn(gindex, *args)
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py", line 523, in tpu_train
self.run_pretrain_routine(model)
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py", line 913, in run_pretrain_routine
self.train()
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 347, in train
self.run_training_epoch()
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 453, in run_training_epoch
self.call_early_stop_callback()
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py", line 793, in call_early_stop_callback
self.early_stop_callback.on_epoch_end(self, self.get_model())
File "/usr/local/lib/python3.6/dist-packages/pytorch_lightning/callbacks/early_stopping.py", line 122, in on_epoch_end
if self.monitor_op(current - self.min_delta, self.best):
RuntimeError: torch_xla/csrc/aten_xla_bridge.cpp:69 : Check failed: xtensor

# Fixed code
def _run_early_stopping_check(self, trainer, pl_module):
    logs = trainer.callback_metrics
    if not self._validate_condition_metric(logs):
        return  # short circuit if metric not present

    current = logs.get(self.monitor)
    if not isinstance(current, torch.Tensor):
        current = torch.tensor(current, device=pl_module.device)

    if self.monitor_op(current - self.min_delta, self.best_score.to(pl_module.device)):
        self.best_score = current
        self.wait_count = 0
    else:
        self.wait_count += 1
        should_stop = self.wait_count >= self.patience

        if bool(should_stop):
            self.stopped_epoch = trainer.current_epoch
            trainer.should_stop = True

    # stop every ddp process if any world process decides to stop
    self._stop_distributed_training(trainer, pl_module)