PinHsuan commited on
Commit
445065e
·
verified ·
1 Parent(s): 562c23b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -38,7 +38,7 @@ css = """
38
  }
39
  """
40
 
41
-
42
  FOLD = 5
43
  model = DualStreamTransformer(n_feat1=25, n_feat2=12, d_model=32).to(DEVICE)
44
  metric_fc = ArcMarginProduct(32, 2).to(DEVICE)
 
38
  }
39
  """
40
 
41
+ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
42
  FOLD = 5
43
  model = DualStreamTransformer(n_feat1=25, n_feat2=12, d_model=32).to(DEVICE)
44
  metric_fc = ArcMarginProduct(32, 2).to(DEVICE)