import torch
import streamlit as st

viz = torch.cuda.memory_summary()
st.write(viz)