Update handler.py

handler.py

@@ -22,6 +22,10 @@ class EndpointHandler:
         )
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         logger.info("Model and tokenizer loaded sucessfully")
+        gpu_memory_allocated = torch.cuda.memory_allocated() / 1024**3  # Convert to GB
+        gpu_memory_reserved = torch.cuda.memory_reserved() / 1024**3  # Convert to GB
+        logger.info(f"GPU memory allocated after model loading: {gpu_memory_allocated:.2f} GB")
+        logger.info(f"GPU memory reserved after model loading: {gpu_memory_reserved:.2f} GB")
         
         # Define the prompt template
         self.prompt_template = """You are an assistant designed to identify whether a user query is malicious.