Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # モデルとトークナイザの初期化
-model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-1b", device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-1b", device_map="cuda"), torch_dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-1b")
 
 # 推論用の関数