diff --git a/containers/chatterbox/server.py b/containers/chatterbox/server.py index ec11344..5d1967c 100644 --- a/containers/chatterbox/server.py +++ b/containers/chatterbox/server.py @@ -18,6 +18,10 @@ elif torch.cuda.is_available(): # ROCm builds expose the GPU via the CUDA API t else: DEVICE = "cpu" +if DEVICE == "cuda": + # avoid MIOpen's broken JIT RNN kernel on some ROCm/GPU combos + torch.backends.cudnn.enabled = False + EXAGGERATION = float(os.environ.get("CHATTERBOX_EXAGGERATION", "0.5")) CFG_WEIGHT = float(os.environ.get("CHATTERBOX_CFG_WEIGHT", "0.5"))