E-Hospital commited on
Commit
91e75bf
1 Parent(s): dcfea26

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +2 -3
main.py CHANGED
@@ -1,9 +1,8 @@
1
  from transformers import AutoTokenizer, AutoModelForCausalLM
2
  import torch
3
- import os
4
- os.environ["CUDA_VISIBLE_DEVICES"]="0"
5
 
6
- device = torch.device("cuda")
 
7
 
8
  model = AutoModelForCausalLM.from_pretrained(
9
  "E-Hospital/open-orca-platypus-2-lora-medical",
 
1
  from transformers import AutoTokenizer, AutoModelForCausalLM
2
  import torch
 
 
3
 
4
+
5
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6
 
7
  model = AutoModelForCausalLM.from_pretrained(
8
  "E-Hospital/open-orca-platypus-2-lora-medical",