Update README.md

README.md

@@ -60,10 +60,17 @@ conda activate lwm_env
 
 Install the necessary packages inside your new environment.
 
+# Install CUDA-enabled Pytorch
+```bash
+pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+```
+# Install other required packages from conda-forge
+```bash
+conda install python numpy pandas matplotlib tqdm -c conda-forge
+```
+# Install DeepMIMOv3 with pip
 ```bash
-conda install -c conda-forge pytorch numpy pandas matplotlib tqdm
 pip install DeepMIMOv3
-
 ```
 
 ---
@@ -180,9 +187,8 @@ else:
 ---
 
 ### 7. **Tokenize and Load the Model**
-```
 
-python
+```python
 from input_preprocess import tokenizer
 from lwm_model import lwm
 import torch