Update main.py

main.py

@@ -1,8 +1,15 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import logging
 
+logger = logging.getLogger()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "CPU") 
+
+num_of_gpus = torch.cuda.device_count()
+logger.info("Getting gpus")
+logger.info(num_of_gpus)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
 
 model = AutoModelForCausalLM.from_pretrained(
     "E-Hospital/open-orca-platypus-2-lora-medical",