Update ovis_image/model/autoencoder.py
Browse files
ovis_image/model/autoencoder.py
CHANGED
|
@@ -363,7 +363,7 @@ class AutoEncoder(nn.Module):
|
|
| 363 |
def load_ae(
|
| 364 |
ckpt_path: str,
|
| 365 |
autoencoder_params: AutoEncoderParams,
|
| 366 |
-
device: str | torch.device = "
|
| 367 |
dtype=torch.bfloat16,
|
| 368 |
random_init=False,
|
| 369 |
) -> AutoEncoder:
|
|
|
|
| 363 |
def load_ae(
|
| 364 |
ckpt_path: str,
|
| 365 |
autoencoder_params: AutoEncoderParams,
|
| 366 |
+
device: str | torch.device = "cpu",
|
| 367 |
dtype=torch.bfloat16,
|
| 368 |
random_init=False,
|
| 369 |
) -> AutoEncoder:
|