1. 初始化代码；

2025-12-30 22:42:47 +08:00
parent 2a75e5acb5
commit aed4f68b0d
11 changed files with 268 additions and 168 deletions
--- a/tools/devices.py
+++ b/tools/devices.py
@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+from functools import cached_property
+
+
+class DeviceManager(object):
+
+    def __init__(self):
+        self.device = self.__get_device()
+
+    @staticmethod
+    def __get_device():
+        if torch.cuda.is_available():
+            return torch.device('cuda:0')
+        if torch.backends.mps.is_available():
+            return torch.device('mps')
+        return torch.device('cpu')