diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py
index 0f6c6ef..080040e 100644
--- a/acestep/pipeline_ace_step.py
+++ b/acestep/pipeline_ace_step.py
@@ -129,8 +129,8 @@ class ACEStepPipeline:
             self.dtype = torch.float16
         if device.type == "mps":
             self.dtype = torch.float32
-        if 'COLAB_FREE' in os.environ and len(os.environ['COLAB_FREE']):
-            self.dtype = torch.float16
+        if 'ACE_PIPELINE_DTYPE' in os.environ and len(os.environ['ACE_PIPELINE_DTYPE']):
+            self.dtype = getattr(torch, os.environ['ACE_PIPELINE_DTYPE'])
         self.device = device
         self.loaded = False
         self.torch_compile = torch_compile
diff --git a/colab_inference.ipynb b/colab_inference.ipynb
index 53c466c..4f787a1 100644
--- a/colab_inference.ipynb
+++ b/colab_inference.ipynb
@@ -47,7 +47,7 @@
         "#!pip uninstall ace-step -y\n",
         "!pip install --upgrade git+https://github.com/craftmine1000/ACE-Step-Colab-Opt.git\n",
         "import os\n",
-        "os.environ['COLAB_FREE'] = 'TRUE'"
+        "os.environ['ACE_PIPELINE_DTYPE'] = 'float16'"
       ]
     },
     {