Improving Machine Learning Frameworks for Catalyst Simulations