handle cpu

README.md

@@ -30,12 +30,18 @@ from model import BertClassifier
 from transformers import PreTrainedTokenizerFast
 import torch
 
-MODEL_PATH = "./pytorch_model.bin"
+# Check if a GPU is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Specify paths
+MODEL_PATH = "./model.pth"
 TOKENIZER_PATH = "./tokenizer.json"
 
+# Load the model with the appropriate map_location
 dzarashield = BertClassifier()
-dzarashield.load_state_dict(torch.load(MODEL_PATH))
+dzarashield.load_state_dict(torch.load(MODEL_PATH, map_location=device))
 
+# Load the tokenizer
 tokenizer = PreTrainedTokenizerFast(tokenizer_file=TOKENIZER_PATH)
 
 ```