gpu

app.py

@@ -13,12 +13,13 @@ goldfish_model = 'goldfish-models/' + model_name
 config = AutoConfig.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
 tokenizer = AutoTokenizer.from_pretrained(goldfish_model, cache_dir=HF_CACHE)
 model = AutoModelForCausalLM.from_pretrained(goldfish_model, config=config, cache_dir=HF_CACHE)
+
 if torch.cuda.is_available():
     model = model.cuda()  # Load onto GPU
 
 
 # Create text generation pipeline
-text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+text_generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
 # Function to generate text
 @spaces.GPU