Initial implementation of saving model weights