Mo-alaa commited on
Commit
a79466d
·
verified ·
1 Parent(s): 4327087

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +3 -1
model.py CHANGED
@@ -1,5 +1,7 @@
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
-
 
 
3
  @torch.no_grad()
4
  def create_model():
5
  tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
+ import torch
3
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
4
+ print(device)
5
  @torch.no_grad()
6
  def create_model():
7
  tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')