Update training dataset to R2E-Gym-Subset.

PiperOrigin-RevId: 904667438

Author: hgao327

examples/deepswe/train_deepswe_nb.py

@@ -421,12 +421,13 @@ def get_lora_model(base_model, model_mesh):
 print("Loading Dataset...")
 
 dataset = load_dataset(
-    "R2E-Gym/R2E-Gym-V1",
+    "R2E-Gym/R2E-Gym-Subset",
     split="train",
     cache_dir=DATASET_CACHE,
     trust_remote_code=True,
 )
 
+
 def transform(entry):
   for k, v in entry.items():
     if isinstance(v, list):