Commit
·
ca78f11
1
Parent(s):
7a96592
Load model to cpu during inference
Browse files- pet_seg_core/webapp.py +1 -1
pet_seg_core/webapp.py
CHANGED
|
@@ -18,7 +18,7 @@ if PetSegWebappConfig.DOWNLOAD_MODEL_WEIGTHS_FROM_GDRIVE:
|
|
| 18 |
PetSegWebappConfig.MODEL_WEIGHTS_GDRIVE_FILE_ID, PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH
|
| 19 |
)
|
| 20 |
|
| 21 |
-
model = UNet.load_from_checkpoint(PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH)
|
| 22 |
model.eval()
|
| 23 |
|
| 24 |
def segment_image(img):
|
|
|
|
| 18 |
PetSegWebappConfig.MODEL_WEIGHTS_GDRIVE_FILE_ID, PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH
|
| 19 |
)
|
| 20 |
|
| 21 |
+
model = UNet.load_from_checkpoint(PetSegWebappConfig.MODEL_WEIGHTS_LOCAL_PATH).to(device)
|
| 22 |
model.eval()
|
| 23 |
|
| 24 |
def segment_image(img):
|