torch 更新

tool/ctx.py

@@ -1,13 +1,17 @@
 import torch
+import spaces
+
 
 class TorchSeedContext:
     def __init__(self, seed):
         self.seed = seed
         self.state = None
 
+    @spaces.GPU
     def __enter__(self):
         self.state = torch.random.get_rng_state()
         torch.manual_seed(self.seed)
 
+    @spaces.GPU
     def __exit__(self, type, value, traceback):
         torch.random.set_rng_state(self.state)