Spaces:
Running
Running
PushkarA07
commited on
Commit
·
47cd4d7
1
Parent(s):
2843087
Update app.py
Browse files
app.py
CHANGED
@@ -355,7 +355,7 @@ from fastai.vision.models.unet import DynamicUnet
|
|
355 |
|
356 |
def build_res_unet(n_input=1, n_output=2, size=256):
|
357 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
358 |
-
body = create_body(resnet18
|
359 |
net_G = DynamicUnet(body, n_output, (size, size)).to(device)
|
360 |
return net_G
|
361 |
|
|
|
355 |
|
356 |
def build_res_unet(n_input=1, n_output=2, size=256):
|
357 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
358 |
+
body = create_body(resnet18, pretrained=True, n_in=n_input, cut=-2)
|
359 |
net_G = DynamicUnet(body, n_output, (size, size)).to(device)
|
360 |
return net_G
|
361 |
|