singhvaibhav924 commited on
Commit
e7186f9
·
1 Parent(s): 4f964a6

minor bug fix

Browse files
Files changed (1) hide show
  1. helper.py +0 -1
helper.py CHANGED
@@ -29,7 +29,6 @@ class KeyphraseExtractionPipeline(TokenClassificationPipeline):
29
  def init_pipeline() :
30
  summarizer_model = AutoModelForCausalLM.from_pretrained(
31
  summarizer_model_name,
32
- device_map="cuda",
33
  torch_dtype=torch.float16,
34
  trust_remote_code=True
35
  )
 
29
  def init_pipeline() :
30
  summarizer_model = AutoModelForCausalLM.from_pretrained(
31
  summarizer_model_name,
 
32
  torch_dtype=torch.float16,
33
  trust_remote_code=True
34
  )