lime-j commited on
Commit
27a537a
·
1 Parent(s): 8989f45
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import numpy as np
2
  import gradio as gr
3
  import torch
@@ -6,7 +7,7 @@ import torch.nn.functional as F
6
  from models.arch.RDnet_ import FullNet_NLP
7
  from models.arch.classifier import PretrainedConvNext
8
  import torchvision.transforms.functional as TF
9
- import spaces
10
 
11
  channels = [64, 128, 256, 512]
12
  layers = [2, 2, 4, 2]
@@ -33,7 +34,7 @@ def align(x1):
33
  def pipe(img):
34
  img=align(img)
35
  image_tensor=TF.to_tensor(img)
36
- image_tensor=image_tensor.unsqueeze(0)
37
  return predict(image_tensor)
38
 
39
  @spaces.GPU(duration=120)
 
1
+ import spaces
2
  import numpy as np
3
  import gradio as gr
4
  import torch
 
7
  from models.arch.RDnet_ import FullNet_NLP
8
  from models.arch.classifier import PretrainedConvNext
9
  import torchvision.transforms.functional as TF
10
+
11
 
12
  channels = [64, 128, 256, 512]
13
  layers = [2, 2, 4, 2]
 
34
  def pipe(img):
35
  img=align(img)
36
  image_tensor=TF.to_tensor(img)
37
+ image_tensor=image_tensor.unsqueeze(0).to('cuda')
38
  return predict(image_tensor)
39
 
40
  @spaces.GPU(duration=120)