We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
minGPT-ddp: AttributeError: 'Trainer' object has no attribute 'dataloader'
When executing examples/distributed/minGPT-ddp/mingpt/main.py

This error is raised when trying to train minGPT.

Python version: main branch

## Possible Solution
113     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):
114         #self.dataloader.sampler.set_epoch(epoch)
115         dataloader.sampler.set_epoch(epoch)


## Steps to Reproduce
Just run main.py

## Failure Logs [if any]

Traceback (most recent call last):
  File "/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/main.py", line 41, in <module>
    main()
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/main.py", line 90, in decorated_main
    _run_hydra(
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py", line 389, in _run_hydra
    _run_app(
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py", line 452, in _run_app
    run_and_report(
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py", line 216, in run_and_report
    raise ex
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py", line 213, in run_and_report
    return func()
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py", line 453, in <lambda>
    lambda: hydra.run(
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/main.py", line 35, in main
    trainer.train()
  File "/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/trainer.py", line 144, in train
    self._run_epoch(epoch, self.train_loader, train=True)
  File "/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/trainer.py", line 114, in _run_epoch
    self.dataloader.sampler.set_epoch(epoch)
AttributeError: 'Trainer' object has no attribute 'dataloader'. Did you mean: 'test_loader'?




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `distributed/minGPT-ddp/mingpt/trainer.py`
Content:
```
1 """
2 Simple training loop; Boilerplate that could apply to any arbitrary neural network,
3 so nothing in this file really has anything to do with GPT specifically.
4 """
5 
6 from dataclasses import dataclass, asdict
7 from collections import OrderedDict
8 from typing import Optional, Any, Dict
9 import os
10 
11 import torch
12 from torch.utils.data import Dataset, DataLoader
13 from torch.nn.parallel import DistributedDataParallel as DDP
14 from torch.utils.data.distributed import DistributedSampler
15 
16 import boto3
17 from urllib.parse import urlparse
18 import fsspec
19 import io
20 
21 @dataclass
22 class TrainerConfig:
23     max_epochs: int = None
24     batch_size: int = None
25     data_loader_workers: int = None
26     grad_norm_clip: float = None
27     snapshot_path: Optional[str] = None
28     save_every: int = None
29     use_amp: bool = None
30 
31 @dataclass
32 class Snapshot:
33     model_state: 'OrderedDict[str, torch.Tensor]'
34     optimizer_state: Dict[str, Any]
35     finished_epoch: int
36 
37 def upload_to_s3(obj, dst):
38     buffer = io.BytesIO()
39     torch.save(obj, buffer)
40     buffer.seek(0)
41     dst = urlparse(dst, allow_fragments=False)
42     boto3.client('s3').upload_fileobj(buffer, dst.netloc, dst.path.lstrip('/'))
43 
44 class Trainer:
45 
46     def __init__(self, trainer_config: TrainerConfig, model, optimizer, train_dataset, test_dataset=None):
47         self.config = trainer_config
48         # set torchrun variables
49         self.local_rank = int(os.environ["LOCAL_RANK"])
50         self.global_rank = int(os.environ["RANK"])  
51         # data stuff
52         self.train_dataset = train_dataset
53         self.train_loader = self._prepare_dataloader(train_dataset)
54         self.test_loader = self._prepare_dataloader(test_dataset) if test_dataset else None
55         # initialize train states
56         self.epochs_run = 0
57         self.model = model.to(self.local_rank)
58         self.optimizer = optimizer        
59         self.save_every = self.config.save_every
60         if self.config.use_amp:
61             self.scaler = torch.cuda.amp.GradScaler()
62         # load snapshot if available. only necessary on the first node.
63         if self.config.snapshot_path is None:
64             self.config.snapshot_path = "snapshot.pt"
65         self._load_snapshot()
66         # wrap with DDP. this step will synch model across all the processes.
67         self.model = DDP(self.model, device_ids=[self.local_rank])
68         
69     def _prepare_dataloader(self, dataset: Dataset):
70         return DataLoader(
71             dataset,
72             batch_size=self.config.batch_size,
73             pin_memory=True,
74             shuffle=False,
75             num_workers=self.config.data_loader_workers,
76             sampler=DistributedSampler(dataset)
77         )
78 
79     def _load_snapshot(self):
80         try:
81             snapshot = fsspec.open(self.config.snapshot_path)
82             with snapshot as f:
83                 snapshot_data = torch.load(f, map_location="cpu")
84         except FileNotFoundError:
85             print("Snapshot not found. Training model from scratch")
86             return 
87 
88         snapshot = Snapshot(**snapshot_data)
89         self.model.load_state_dict(snapshot.model_state)
90         self.optimizer.load_state_dict(snapshot.optimizer_state)
91         self.epochs_run = snapshot.finished_epoch
92         print(f"Resuming training from snapshot at Epoch {self.epochs_run}")
93 
94 
95     def _run_batch(self, source, targets, train: bool = True) -> float:
96         with torch.set_grad_enabled(train), torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=(self.config.use_amp)):
97             _, loss = self.model(source, targets)
98         
99         if train:
100             self.optimizer.zero_grad(set_to_none=True)
101             if self.config.use_amp: 
102                 self.scaler.scale(loss).backward()
103                 torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)
104                 self.scaler.step(self.optimizer)
105                 self.scaler.update()
106             else:
107                 loss.backward()
108                 torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)
109                 self.optimizer.step()
110         
111         return loss.item()
112 
113     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):
114         self.dataloader.sampler.set_epoch(epoch)
115         for iter, (source, targets) in enumerate(dataloader):
116             step_type = "Train" if train else "Eval"
117             source = source.to(self.local_rank)
118             targets = targets.to(self.local_rank)
119             batch_loss = self._run_batch(source, targets, train)
120             if iter % 100 == 0:
121                 print(f"[GPU{self.global_rank}] Epoch {epoch} | Iter {iter} | {step_type} Loss {batch_loss:.5f}")
122 
123     def _save_snapshot(self, epoch):
124         # capture snapshot
125         model = self.model
126         raw_model = model.module if hasattr(model, "module") else model
127         snapshot = Snapshot(
128             model_state=raw_model.state_dict(),
129             optimizer_state=self.optimizer.state_dict(),
130             finished_epoch=epoch
131         )
132         # save snapshot
133         snapshot = asdict(snapshot)
134         if self.config.snapshot_path.startswith("s3://"):
135             upload_to_s3(snapshot, self.config.snapshot_path)
136         else:
137             torch.save(snapshot, self.config.snapshot_path)
138             
139         print(f"Snapshot saved at epoch {epoch}")
140 
141     def train(self):
142         for epoch in range(self.epochs_run, self.config.max_epochs):
143             epoch += 1
144             self._run_epoch(epoch, self.train_loader, train=True)
145             if self.local_rank == 0 and epoch % self.save_every == 0:
146                 self._save_snapshot(epoch)
147             # eval run
148             if self.test_loader:
149                 self._run_epoch(epoch, self.test_loader, train=False)
150
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/distributed/minGPT-ddp/mingpt/trainer.py b/distributed/minGPT-ddp/mingpt/trainer.py
--- a/distributed/minGPT-ddp/mingpt/trainer.py
+++ b/distributed/minGPT-ddp/mingpt/trainer.py
@@ -111,7 +111,7 @@
         return loss.item()
 
     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):
-        self.dataloader.sampler.set_epoch(epoch)
+        dataloader.sampler.set_epoch(epoch)
         for iter, (source, targets) in enumerate(dataloader):
             step_type = "Train" if train else "Eval"
             source = source.to(self.local_rank)

{"golden_diff": "diff --git a/distributed/minGPT-ddp/mingpt/trainer.py b/distributed/minGPT-ddp/mingpt/trainer.py\n--- a/distributed/minGPT-ddp/mingpt/trainer.py\n+++ b/distributed/minGPT-ddp/mingpt/trainer.py\n@@ -111,7 +111,7 @@\n         return loss.item()\n \n     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):\n-        self.dataloader.sampler.set_epoch(epoch)\n+        dataloader.sampler.set_epoch(epoch)\n         for iter, (source, targets) in enumerate(dataloader):\n             step_type = \"Train\" if train else \"Eval\"\n             source = source.to(self.local_rank)\n", "issue": "minGPT-ddp: AttributeError: 'Trainer' object has no attribute 'dataloader'\nWhen executing examples/distributed/minGPT-ddp/mingpt/main.py\r\n\r\nThis error is raised when trying to train minGPT.\r\n\r\nPython version: main branch\r\n\r\n## Possible Solution\r\n113     def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):\r\n114         #self.dataloader.sampler.set_epoch(epoch)\r\n115         dataloader.sampler.set_epoch(epoch)\r\n\r\n\r\n## Steps to Reproduce\r\nJust run main.py\r\n\r\n## Failure Logs [if any]\r\n\r\nTraceback (most recent call last):\r\n  File \"/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/main.py\", line 41, in <module>\r\n    main()\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/main.py\", line 90, in decorated_main\r\n    _run_hydra(\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 389, in _run_hydra\r\n    _run_app(\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 452, in _run_app\r\n    run_and_report(\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 216, in run_and_report\r\n    raise ex\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 213, in run_and_report\r\n    return func()\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/utils.py\", line 453, in <lambda>\r\n    lambda: hydra.run(\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/_internal/hydra.py\", line 132, in run\r\n    _ = ret.return_value\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/core/utils.py\", line 260, in return_value\r\n    raise self._return_value\r\n  File \"/home/users/fmansouri/.local/lib/python3.10/site-packages/hydra/core/utils.py\", line 186, in run_job\r\n    ret.return_value = task_function(task_cfg)\r\n  File \"/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/main.py\", line 35, in main\r\n    trainer.train()\r\n  File \"/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/trainer.py\", line 144, in train\r\n    self._run_epoch(epoch, self.train_loader, train=True)\r\n  File \"/mnt/tier1/project/lxp/fmansouri/pytorch/examples/distributed/minGPT-ddp/mingpt/trainer.py\", line 114, in _run_epoch\r\n    self.dataloader.sampler.set_epoch(epoch)\r\nAttributeError: 'Trainer' object has no attribute 'dataloader'. Did you mean: 'test_loader'?\r\n\r\n\r\n\n", "before_files": [{"content": "\"\"\"\nSimple training loop; Boilerplate that could apply to any arbitrary neural network,\nso nothing in this file really has anything to do with GPT specifically.\n\"\"\"\n\nfrom dataclasses import dataclass, asdict\nfrom collections import OrderedDict\nfrom typing import Optional, Any, Dict\nimport os\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data.distributed import DistributedSampler\n\nimport boto3\nfrom urllib.parse import urlparse\nimport fsspec\nimport io\n\n@dataclass\nclass TrainerConfig:\n    max_epochs: int = None\n    batch_size: int = None\n    data_loader_workers: int = None\n    grad_norm_clip: float = None\n    snapshot_path: Optional[str] = None\n    save_every: int = None\n    use_amp: bool = None\n\n@dataclass\nclass Snapshot:\n    model_state: 'OrderedDict[str, torch.Tensor]'\n    optimizer_state: Dict[str, Any]\n    finished_epoch: int\n\ndef upload_to_s3(obj, dst):\n    buffer = io.BytesIO()\n    torch.save(obj, buffer)\n    buffer.seek(0)\n    dst = urlparse(dst, allow_fragments=False)\n    boto3.client('s3').upload_fileobj(buffer, dst.netloc, dst.path.lstrip('/'))\n\nclass Trainer:\n\n    def __init__(self, trainer_config: TrainerConfig, model, optimizer, train_dataset, test_dataset=None):\n        self.config = trainer_config\n        # set torchrun variables\n        self.local_rank = int(os.environ[\"LOCAL_RANK\"])\n        self.global_rank = int(os.environ[\"RANK\"])  \n        # data stuff\n        self.train_dataset = train_dataset\n        self.train_loader = self._prepare_dataloader(train_dataset)\n        self.test_loader = self._prepare_dataloader(test_dataset) if test_dataset else None\n        # initialize train states\n        self.epochs_run = 0\n        self.model = model.to(self.local_rank)\n        self.optimizer = optimizer        \n        self.save_every = self.config.save_every\n        if self.config.use_amp:\n            self.scaler = torch.cuda.amp.GradScaler()\n        # load snapshot if available. only necessary on the first node.\n        if self.config.snapshot_path is None:\n            self.config.snapshot_path = \"snapshot.pt\"\n        self._load_snapshot()\n        # wrap with DDP. this step will synch model across all the processes.\n        self.model = DDP(self.model, device_ids=[self.local_rank])\n        \n    def _prepare_dataloader(self, dataset: Dataset):\n        return DataLoader(\n            dataset,\n            batch_size=self.config.batch_size,\n            pin_memory=True,\n            shuffle=False,\n            num_workers=self.config.data_loader_workers,\n            sampler=DistributedSampler(dataset)\n        )\n\n    def _load_snapshot(self):\n        try:\n            snapshot = fsspec.open(self.config.snapshot_path)\n            with snapshot as f:\n                snapshot_data = torch.load(f, map_location=\"cpu\")\n        except FileNotFoundError:\n            print(\"Snapshot not found. Training model from scratch\")\n            return \n\n        snapshot = Snapshot(**snapshot_data)\n        self.model.load_state_dict(snapshot.model_state)\n        self.optimizer.load_state_dict(snapshot.optimizer_state)\n        self.epochs_run = snapshot.finished_epoch\n        print(f\"Resuming training from snapshot at Epoch {self.epochs_run}\")\n\n\n    def _run_batch(self, source, targets, train: bool = True) -> float:\n        with torch.set_grad_enabled(train), torch.amp.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=(self.config.use_amp)):\n            _, loss = self.model(source, targets)\n        \n        if train:\n            self.optimizer.zero_grad(set_to_none=True)\n            if self.config.use_amp: \n                self.scaler.scale(loss).backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n            else:\n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)\n                self.optimizer.step()\n        \n        return loss.item()\n\n    def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):\n        self.dataloader.sampler.set_epoch(epoch)\n        for iter, (source, targets) in enumerate(dataloader):\n            step_type = \"Train\" if train else \"Eval\"\n            source = source.to(self.local_rank)\n            targets = targets.to(self.local_rank)\n            batch_loss = self._run_batch(source, targets, train)\n            if iter % 100 == 0:\n                print(f\"[GPU{self.global_rank}] Epoch {epoch} | Iter {iter} | {step_type} Loss {batch_loss:.5f}\")\n\n    def _save_snapshot(self, epoch):\n        # capture snapshot\n        model = self.model\n        raw_model = model.module if hasattr(model, \"module\") else model\n        snapshot = Snapshot(\n            model_state=raw_model.state_dict(),\n            optimizer_state=self.optimizer.state_dict(),\n            finished_epoch=epoch\n        )\n        # save snapshot\n        snapshot = asdict(snapshot)\n        if self.config.snapshot_path.startswith(\"s3://\"):\n            upload_to_s3(snapshot, self.config.snapshot_path)\n        else:\n            torch.save(snapshot, self.config.snapshot_path)\n            \n        print(f\"Snapshot saved at epoch {epoch}\")\n\n    def train(self):\n        for epoch in range(self.epochs_run, self.config.max_epochs):\n            epoch += 1\n            self._run_epoch(epoch, self.train_loader, train=True)\n            if self.local_rank == 0 and epoch % self.save_every == 0:\n                self._save_snapshot(epoch)\n            # eval run\n            if self.test_loader:\n                self._run_epoch(epoch, self.test_loader, train=False)\n", "path": "distributed/minGPT-ddp/mingpt/trainer.py"}], "after_files": [{"content": "\"\"\"\nSimple training loop; Boilerplate that could apply to any arbitrary neural network,\nso nothing in this file really has anything to do with GPT specifically.\n\"\"\"\n\nfrom dataclasses import dataclass, asdict\nfrom collections import OrderedDict\nfrom typing import Optional, Any, Dict\nimport os\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data.distributed import DistributedSampler\n\nimport boto3\nfrom urllib.parse import urlparse\nimport fsspec\nimport io\n\n@dataclass\nclass TrainerConfig:\n    max_epochs: int = None\n    batch_size: int = None\n    data_loader_workers: int = None\n    grad_norm_clip: float = None\n    snapshot_path: Optional[str] = None\n    save_every: int = None\n    use_amp: bool = None\n\n@dataclass\nclass Snapshot:\n    model_state: 'OrderedDict[str, torch.Tensor]'\n    optimizer_state: Dict[str, Any]\n    finished_epoch: int\n\ndef upload_to_s3(obj, dst):\n    buffer = io.BytesIO()\n    torch.save(obj, buffer)\n    buffer.seek(0)\n    dst = urlparse(dst, allow_fragments=False)\n    boto3.client('s3').upload_fileobj(buffer, dst.netloc, dst.path.lstrip('/'))\n\nclass Trainer:\n\n    def __init__(self, trainer_config: TrainerConfig, model, optimizer, train_dataset, test_dataset=None):\n        self.config = trainer_config\n        # set torchrun variables\n        self.local_rank = int(os.environ[\"LOCAL_RANK\"])\n        self.global_rank = int(os.environ[\"RANK\"])  \n        # data stuff\n        self.train_dataset = train_dataset\n        self.train_loader = self._prepare_dataloader(train_dataset)\n        self.test_loader = self._prepare_dataloader(test_dataset) if test_dataset else None\n        # initialize train states\n        self.epochs_run = 0\n        self.model = model.to(self.local_rank)\n        self.optimizer = optimizer        \n        self.save_every = self.config.save_every\n        if self.config.use_amp:\n            self.scaler = torch.cuda.amp.GradScaler()\n        # load snapshot if available. only necessary on the first node.\n        if self.config.snapshot_path is None:\n            self.config.snapshot_path = \"snapshot.pt\"\n        self._load_snapshot()\n        # wrap with DDP. this step will synch model across all the processes.\n        self.model = DDP(self.model, device_ids=[self.local_rank])\n        \n    def _prepare_dataloader(self, dataset: Dataset):\n        return DataLoader(\n            dataset,\n            batch_size=self.config.batch_size,\n            pin_memory=True,\n            shuffle=False,\n            num_workers=self.config.data_loader_workers,\n            sampler=DistributedSampler(dataset)\n        )\n\n    def _load_snapshot(self):\n        try:\n            snapshot = fsspec.open(self.config.snapshot_path)\n            with snapshot as f:\n                snapshot_data = torch.load(f, map_location=\"cpu\")\n        except FileNotFoundError:\n            print(\"Snapshot not found. Training model from scratch\")\n            return \n\n        snapshot = Snapshot(**snapshot_data)\n        self.model.load_state_dict(snapshot.model_state)\n        self.optimizer.load_state_dict(snapshot.optimizer_state)\n        self.epochs_run = snapshot.finished_epoch\n        print(f\"Resuming training from snapshot at Epoch {self.epochs_run}\")\n\n\n    def _run_batch(self, source, targets, train: bool = True) -> float:\n        with torch.set_grad_enabled(train), torch.amp.autocast(device_type=\"cuda\", dtype=torch.float16, enabled=(self.config.use_amp)):\n            _, loss = self.model(source, targets)\n        \n        if train:\n            self.optimizer.zero_grad(set_to_none=True)\n            if self.config.use_amp: \n                self.scaler.scale(loss).backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n            else:\n                loss.backward()\n                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.grad_norm_clip)\n                self.optimizer.step()\n        \n        return loss.item()\n\n    def _run_epoch(self, epoch: int, dataloader: DataLoader, train: bool = True):\n        dataloader.sampler.set_epoch(epoch)\n        for iter, (source, targets) in enumerate(dataloader):\n            step_type = \"Train\" if train else \"Eval\"\n            source = source.to(self.local_rank)\n            targets = targets.to(self.local_rank)\n            batch_loss = self._run_batch(source, targets, train)\n            if iter % 100 == 0:\n                print(f\"[GPU{self.global_rank}] Epoch {epoch} | Iter {iter} | {step_type} Loss {batch_loss:.5f}\")\n\n    def _save_snapshot(self, epoch):\n        # capture snapshot\n        model = self.model\n        raw_model = model.module if hasattr(model, \"module\") else model\n        snapshot = Snapshot(\n            model_state=raw_model.state_dict(),\n            optimizer_state=self.optimizer.state_dict(),\n            finished_epoch=epoch\n        )\n        # save snapshot\n        snapshot = asdict(snapshot)\n        if self.config.snapshot_path.startswith(\"s3://\"):\n            upload_to_s3(snapshot, self.config.snapshot_path)\n        else:\n            torch.save(snapshot, self.config.snapshot_path)\n            \n        print(f\"Snapshot saved at epoch {epoch}\")\n\n    def train(self):\n        for epoch in range(self.epochs_run, self.config.max_epochs):\n            epoch += 1\n            self._run_epoch(epoch, self.train_loader, train=True)\n            if self.local_rank == 0 and epoch % self.save_every == 0:\n                self._save_snapshot(epoch)\n            # eval run\n            if self.test_loader:\n                self._run_epoch(epoch, self.test_loader, train=False)\n", "path": "distributed/minGPT-ddp/mingpt/trainer.py"}]}