This commit is contained in:
语帆 2024-02-28 15:23:07 +08:00
parent eb92e79fb9
commit debafeac37

View File

@ -141,7 +141,7 @@ class AutoModel:
kwargs = download_model(**kwargs)
set_all_random_seed(kwargs.get("seed", 0))
pdb.set_trace()
device = kwargs.get("device", "cuda")
if not torch.cuda.is_available() or kwargs.get("ngpu", 1) == 0:
device = "cpu"