RWKV-Runner/backend-python/utils/torch.py
2023-05-07 17:27:54 +08:00

27 lines
731 B
Python

import os
import sysconfig
def set_torch():
torch_path = os.path.join(sysconfig.get_paths()["purelib"], "torch\\lib")
paths = os.environ.get("PATH", "")
if os.path.exists(torch_path):
print(f"torch found: {torch_path}")
if torch_path in paths:
print("torch already set")
else:
print("run:")
os.environ['PATH'] = paths + os.pathsep + torch_path + os.pathsep
print(f'set Path={paths + os.pathsep + torch_path + os.pathsep}')
else:
print("torch not found")
def torch_gc():
import torch
if torch.cuda.is_available():
with torch.cuda.device(0):
torch.cuda.empty_cache()
torch.cuda.ipc_collect()