diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index f724650f3..5f89dd671 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -146,7 +146,7 @@ class AutoModel: set_all_random_seed(kwargs.get("seed", 0)) device = kwargs.get("device", "cuda") - if not torch.cuda.is_available() or kwargs.get("ngpu", 0): + if not torch.cuda.is_available() or kwargs.get("ngpu", 0) == 0: device = "cpu" kwargs["batch_size"] = 1 kwargs["device"] = device