loading VAE from safetensors

ciCciC · Dec 4, 2023 · bbf241b · bbf241b
1 parent 5e1ba96
commit bbf241b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/app/models/compressor.py b/app/models/compressor.py
@@ -25,7 +25,9 @@ def get_latent_mu_size(self):
 
     def load_model(self):
         self._model = (AutoencoderTiny
-                       .from_pretrained(self.model_id, torch_dtype=self.d_type)
+                       .from_pretrained(self.model_id,
+                                        torch_dtype=self.d_type,
+                                        use_safetensors=True)
                        .to(self._device))
 
     def get_model(self):