pytorch / xla

Enabling PyTorch on XLA Devices (e.g. Google TPU)
https://pytorch.org/xla
Other
2.49k stars 480 forks source link

Update and short circuit the Neuron initialization #8041

Closed rpsilva-aws closed 1 month ago

rpsilva-aws commented 2 months ago

Currently, we re-initialize the environment for Neuron when torch-neuronx is present. In this change, we decouple the initialization when torch-neuronx is present. In the meantime, we also short circuit the returning value if libneuronxla is not present.