fix: fix type annotation.

Author: GdoongMathew

src/lightning/pytorch/accelerators/accelerator.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC
-from typing import Any, Optional
+from typing import Any
 
 import lightning.pytorch as pl
 from lightning.fabric.accelerators.accelerator import Accelerator as _Accelerator
@@ -47,6 +47,6 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]:
         raise NotImplementedError
 
     @classmethod
-    def device_name(cls, device: Optional = None) -> str:
+    def device_name(cls, device: _DEVICE = None) -> str:
         """Get the device name for a given device."""
         return str(cls.is_available())

src/lightning/pytorch/accelerators/cuda.py

@@ -114,7 +114,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No
         )
 
     @classmethod
-    def device_name(cls, device: Optional[torch.types.Device] = None) -> str:
+    def device_name(cls, device: _DEVICE = None) -> str:
         if not cls.is_available():
             return "False"
         return torch.cuda.get_device_name(device)

src/lightning/pytorch/accelerators/mps.py

@@ -88,7 +88,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No
         )
 
     @classmethod
-    def device_name(cls, device: Optional = None) -> str:
+    def device_name(cls, device: _DEVICE = None) -> str:
         # todo: implement a better way to get the device name
         available = cls.is_available()
         gpu_type = " (mps)" if available else ""

src/lightning/pytorch/accelerators/xla.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Any
 
 from typing_extensions import override
 
@@ -56,7 +56,7 @@ def register_accelerators(cls, accelerator_registry: _AcceleratorRegistry) -> No
         accelerator_registry.register("tpu", cls, description=cls.__name__)
 
     @classmethod
-    def device_name(cls, device: Optional = None) -> str:
+    def device_name(cls, device: _DEVICE = None) -> str:
         is_available = cls.is_available()
         if not is_available:
             return str(is_available)