Commit
·
9a4a6ce
1
Parent(s):
5e88f71
change directory absolute at unet model directory
Browse files
app.py
CHANGED
|
@@ -352,7 +352,7 @@ IMG_SIZE = 128 # Resize dimension for the input image
|
|
| 352 |
@st.cache_resource
|
| 353 |
def load_model():
|
| 354 |
model = UNet(n_channels=3, n_classes=32) # Adjust according to your model setup
|
| 355 |
-
model.load_state_dict(torch.load("
|
| 356 |
model.eval()
|
| 357 |
return model
|
| 358 |
|
|
|
|
| 352 |
@st.cache_resource
|
| 353 |
def load_model():
|
| 354 |
model = UNet(n_channels=3, n_classes=32) # Adjust according to your model setup
|
| 355 |
+
model.load_state_dict(torch.load("./model/unet/checkpoint_epoch5.pth", map_location="cpu", weights_only=True), strict=False)
|
| 356 |
model.eval()
|
| 357 |
return model
|
| 358 |
|