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 |
|