wing-nus commited on
Commit
516a961
1 Parent(s): 61f57fc

Update dataset_extraction.py

Browse files
Files changed (1) hide show
  1. dataset_extraction.py +1 -1
dataset_extraction.py CHANGED
@@ -4,7 +4,7 @@ import nltk
4
  from SciAssist import DatasetExtraction
5
 
6
  device = "gpu" if torch.cuda.is_available() else "cpu"
7
- de_pipeline = DatasetExtraction(os_name="nt")
8
 
9
 
10
  def de_for_str(input):
 
4
  from SciAssist import DatasetExtraction
5
 
6
  device = "gpu" if torch.cuda.is_available() else "cpu"
7
+ de_pipeline = DatasetExtraction(os_name="nt", device=device)
8
 
9
 
10
  def de_for_str(input):