Spaces:
Running
Running
Test with checkpoint 24
Browse files
app.py
CHANGED
@@ -17,8 +17,8 @@ from skimage.metrics import structural_similarity as ssim_metric
|
|
17 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
18 |
|
19 |
from model import ViTUNetColorizer
|
20 |
-
CKPT = "checkpoints/checkpoint_epoch_015_20250808_154437.pt"
|
21 |
-
|
22 |
model = None
|
23 |
if os.path.exists(CKPT):
|
24 |
model = ViTUNetColorizer(vit_model_name="vit_tiny_patch16_224").to(device)
|
|
|
17 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
18 |
|
19 |
from model import ViTUNetColorizer
|
20 |
+
# CKPT = "checkpoints/checkpoint_epoch_015_20250808_154437.pt"
|
21 |
+
CKPT = "checkpoints/checkpoint_epoch_024_20250808_200347.pt"
|
22 |
model = None
|
23 |
if os.path.exists(CKPT):
|
24 |
model = ViTUNetColorizer(vit_model_name="vit_tiny_patch16_224").to(device)
|