save

harrisonvanderbyl · Mar 18, 2023 · 542f51c · 542f51c
1 parent 2cd56b7
commit 542f51c
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/src/rwkvstic/agnostic/backends/opt.py b/src/rwkvstic/agnostic/backends/opt.py
@@ -12,14 +12,12 @@
 import os
 current_path = os.path.dirname(os.path.abspath(__file__))
 
-# method = torch.jit.script_method
-# module = torch.jit.ScriptModule
-# script = torch.jit.script
-method = lambda x:x #.script_method
-module = torch.nn.Module #.ScriptModule
-script = lambda x:x #.script
+method = torch.jit.script_method
+module = torch.jit.ScriptModule
+script = torch.jit.script
 with torch.no_grad():
     def OptRWKV(path, jit=True, export=False, **kwargs):
+
 
         device = "cuda"