Update README.md

README.md

@@ -27,7 +27,13 @@ metrics:
 
 ## Setup
 
-You can run the smashed model by **(1)** installing and importing the `pruna-engine` (version 0.2.5) package with the [Pypi instructions](https://pypi.org/project/pruna-engine/),**(2)** downloading the model files at `model_path`. This can be done using huggingface with this repository name or with manual downloading, and **(3)** loading the model, **(4)** running the model. You can achieve this by running the following code:
+You can run the smashed model on a A100 GPU by:
+1. Installing and importing the `pruna-engine` (version 0.2.6) package. Use `pip install pruna --extra-index-url https://pypi.nvidia.com --extra-index-url https://pypi.ngc.nvidia.com` for installation. See [Pypi](https://pypi.org/project/pruna-engine/) for detailed on the package.
+2. Downloading the model files at `model_path`. This can be done using huggingface with this repository name or with manual downloading.
+3. Loading the model
+4. Running the model. 
+
+You can achieve this by running the following code:
 
 ```python
 from transformers.utils.hub import cached_file
@@ -36,7 +42,7 @@ from pruna_engine.PrunaModel import PrunaModel  # Step (1): install and import `
 ...
 model_path = cached_file("PrunaAI/REPO", "model")  # Step (2): download the model files at `model_path`.
 smashed_model = PrunaModel.load_model(model_path)  # Step (3): load the model.
-y = smashed_model(x)  # Step (4): run the model.
+y = smashed_model(prompt="a photo of an astronaut riding a horse on mars", image_height=1024, image_width=1024)[0]  # Step (4): run the model.
 ```
 
 ## Configurations