GpuPlugin -> CudaPlugin

pytorch · Jan 18, 2024 · cfb60fa · cfb60fa
1 parent 2bc097e
commit cfb60fa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/plugins/cuda/README.md b/plugins/cuda/README.md
@@ -29,7 +29,7 @@ import torch_xla.runtime as xr
 
 # Use dynamic plugin instead of built-in CUDA support
 plugins.use_dynamic_plugins()
-plugins.register_plugin('CUDA', torch_xla_cuda_plugin.GpuPlugin())
+plugins.register_plugin('CUDA', torch_xla_cuda_plugin.CudaPlugin())
 xr.set_device_type('CUDA')
 
 print(xm.xla_device())

diff --git a/plugins/cuda/pyproject.toml b/plugins/cuda/pyproject.toml
@@ -15,4 +15,4 @@ requires-python = ">=3.8"
 torch_xla_cuda_plugin = ["lib/*.so"]
 
 [project.entry-points."torch_xla.plugins"]
-gpu = "torch_xla_cuda_plugin:GpuPlugin"
+cuda = "torch_xla_cuda_plugin:CudaPlugin"