Update app.py

app.py

@@ -483,18 +483,22 @@ with tabs[0]:
                 # Format details
                 details = ""
                 if show_details:
+                    gpu_memory_used = f"{float(final_perf.get('gpu_memory_used', 0)):.2f} MB" if DEVICE == "cuda" else "N/A"
+                    gpu_memory_cached = f"{float(final_perf.get('gpu_memory_cached', 0)):.2f} MB" if DEVICE == "cuda" else "N/A"
+
                     details = f"""
                         3-Class Probability: {result['prob3class']:.2f}
                         3-Class Predicted Label: {['NEI', 'SUPPORTED', 'REFUTED'][result['pred_tc']]}
                         2-Class Probability: {result['prob2class']:.2f}
                         2-Class Predicted Label: {['SUPPORTED', 'REFUTED'][result['pred_bc']] if isinstance(result['pred_bc'], int) and result['pred_tc'] != 0 else 'Not used'}
-                        
+
                         Performance Metrics:
-                        - GPU Memory Used: {final_perf.get('gpu_memory_used', 'N/A'):.2f} MB
-                        - GPU Memory Cached: {final_perf.get('gpu_memory_cached', 'N/A'):.2f} MB
+                        - GPU Memory Used: {gpu_memory_used}
+                        - GPU Memory Cached: {gpu_memory_cached}
                         - CPU Usage: {final_perf['cpu_percent']}%
                         - Memory Usage: {final_perf['memory_percent']}%
                     """
+
                 
                 # Store result with performance metrics
                 st.session_state.latest_result = {
@@ -524,6 +528,12 @@ with tabs[0]:
                     "NEI": "verdict-nei"
                 }.get(res['verdict'], "")
                 
+                gpu_memory_text = (
+                    f"<li>GPU Memory: {float(res['performance_metrics'].get('gpu_memory_used', 0)):.2f} MB</li>"
+                    if DEVICE == "cuda"
+                    else "<li>GPU Memory: N/A</li>"
+                )
+
                 st.markdown(f"""
                     <div class="result-box">
                         <h3>Verification Results</h3>
@@ -540,7 +550,7 @@ with tabs[0]:
                             <ul>
                                 <li>CPU: {res['performance_metrics']['cpu_percent']}%</li>
                                 <li>RAM: {res['performance_metrics']['memory_percent']}%</li>
-                                {f"<li>GPU Memory: {res['performance_metrics'].get('gpu_memory_used', 'N/A'):.2f} MB</li>" if DEVICE == "cuda" else ""}
+                                {gpu_memory_text}
                             </ul>
                         </div>
                         {f"<div class='code-block'><pre>{res['details']}</pre></div>" if show_details else ""}
@@ -557,7 +567,7 @@ Details: {res['details']}
 Performance:
 - CPU: {res['performance_metrics']['cpu_percent']}%
 - RAM: {res['performance_metrics']['memory_percent']}%
-{f"- GPU Memory: {res['performance_metrics'].get('gpu_memory_used', 'N/A'):.2f} MB" if DEVICE == "cuda" else ""}
+- GPU Memory: {f"{float(res['performance_metrics'].get('gpu_memory_used', 0)):.2f} MB" if DEVICE == "cuda" else "N/A"}
 """
                 st.download_button(
                     "📥 Download Results",