Update README.md

README.md

@@ -12,7 +12,7 @@ This network uses the optimal transport flow matching objective outlined [Flow M
 xATGLU Layers are used instead of linears for entry into the transformer MLP layer [Expanded Gating Ranges
 Improve Activation Functions](https://arxiv.org/pdf/2405.20768)
 
-```python train.py``` will train a new image network on the provided dataset. 
+```python train.py``` will train a new image network on the provided dataset (Currently the dataset is being fully rammed into GPU and is defined in the preload_dataset function)
 
 ```python test_sample.py step_1799.safetensors``` Where step_1799.safetensors is the desired model to test inference on. This will always generate a sample grid of 16x16 images.