Spaces:
Runtime error
Runtime error
add model
Browse files- App_main.py +2 -1
App_main.py
CHANGED
@@ -2,7 +2,7 @@ from collections import OrderedDict
|
|
2 |
|
3 |
import gradio as gr
|
4 |
import os
|
5 |
-
|
6 |
import torch
|
7 |
from torchvision import transforms
|
8 |
from torchvision.transforms import InterpolationMode
|
@@ -55,6 +55,7 @@ def seg_buildings(Image, Checkpoint):
|
|
55 |
sample = construct_sample(Image, mean, std)
|
56 |
model = build_model(checkpoint)
|
57 |
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
|
|
58 |
|
59 |
model = model.to(device)
|
60 |
model.eval()
|
|
|
2 |
|
3 |
import gradio as gr
|
4 |
import os
|
5 |
+
os.system('pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113')
|
6 |
import torch
|
7 |
from torchvision import transforms
|
8 |
from torchvision.transforms import InterpolationMode
|
|
|
55 |
sample = construct_sample(Image, mean, std)
|
56 |
model = build_model(checkpoint)
|
57 |
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
58 |
+
print('Use: ', device)
|
59 |
|
60 |
model = model.to(device)
|
61 |
model.eval()
|