Update main.py

main.py

@@ -118,6 +118,9 @@ def clean_memory():
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        # Clear GPU cache
+        torch.cuda.empty_cache()
+        logger.info(f"Memory allocated after clearing cache: {torch.cuda.memory_allocated()} bytes")
         logger.info("Memory cleaned")
 
 def predict(image):