Utilities used by other modules
xla_imported
is a utility method that is used to check if the torch_xla
module has been successfully imported.
One of the problems we have hit testing different models and transforms is that sometimes it is slower on TPUs compared to running on CPUs, but this happens because we hit operations on Pytorch XLA that are only handled by the CPU and not by the accelerator.
print_aten_ops
calls directly some pytorch metrics which outputs to stdout
, so the only way to get that info is capture it.
#test that torch_xla has been imported on colab
assert xla_imported()