hiascend commited on
Commit
376c53e
·
verified ·
1 Parent(s): c311b2d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -2,6 +2,8 @@ import gradio as gr
2
  import spaces
3
  import torch
4
 
 
 
5
 
6
  class LinearModel(torch.nn.Module):
7
  def __init__(self):
@@ -15,7 +17,6 @@ class LinearModel(torch.nn.Module):
15
 
16
 
17
  def process():
18
- torch.set_default_device('cuda')
19
  print('cuda avaliable: ' + str(torch.cuda.is_available()))
20
  print('cuda device count: ' + str(torch.cuda.device_count()))
21
 
 
2
  import spaces
3
  import torch
4
 
5
+ torch.set_default_device('cuda')
6
+
7
 
8
  class LinearModel(torch.nn.Module):
9
  def __init__(self):
 
17
 
18
 
19
  def process():
 
20
  print('cuda avaliable: ' + str(torch.cuda.is_available()))
21
  print('cuda device count: ' + str(torch.cuda.device_count()))
22