diff --git a/scripts/streamlit/depth2img.py b/scripts/streamlit/depth2img.py
index 7f80223..b856481 100644
--- a/scripts/streamlit/depth2img.py
+++ b/scripts/streamlit/depth2img.py
@@ -62,7 +62,7 @@ def paint(sampler, image, prompt, t_enc, seed, scale, num_samples=1, callback=No
     wm_encoder.set_watermark('bytes', wm.encode('utf-8'))
 
     with torch.no_grad(),\
-            torch.autocast("cuda"):
+            torch.onnx.set_training(model, False):
         batch = make_batch_sd(image, txt=prompt, device=device, num_samples=num_samples)
         z = model.get_first_stage_encoding(model.encode_first_stage(batch[model.first_stage_key]))  # move to latent space
         c = model.cond_stage_model.encode(batch["txt"])