prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  TrainConfig parameter "epochs"
  **TrainConfig parameter "epochs" doesn't have effect.**
  After changing the number of epochs=1 to epochs=100. The worker still do only 1 epoch.
  
  ```
  train_config = sy.TrainConfig(
          model=traced_model,
          loss_fn=loss_fn,
          batch_size=batch_size,
          shuffle=True,
          #max_nr_batches=max_nr_batches,
          epochs=100,
          lr=lr,
      )
  ```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `syft/federated/federated_client.py`
  Content:
  ```
  1 import torch as th
  2 from torch.utils.data import BatchSampler, RandomSampler, SequentialSampler
  3 
  4 from syft.generic import ObjectStorage
  5 from syft.federated.train_config import TrainConfig
  6 
  7 
  8 class FederatedClient(ObjectStorage):
  9     """A Client able to execute federated learning in local datasets."""
  10 
  11     def __init__(self, datasets=None):
  12         super().__init__()
  13         self.datasets = datasets if datasets is not None else dict()
  14         self.optimizer = None
  15         self.train_config = None
  16 
  17     def add_dataset(self, dataset, key: str):
  18         self.datasets[key] = dataset
  19 
  20     def remove_dataset(self, key: str):
  21         if key in self.datasets:
  22             del self.datasets[key]
  23 
  24     def set_obj(self, obj: object):
  25         """Registers objects checking if which objects it should cache.
  26 
  27         Args:
  28             obj: An object to be registered.
  29         """
  30         if isinstance(obj, TrainConfig):
  31             self.train_config = obj
  32             self.optimizer = None
  33         else:
  34             super().set_obj(obj)
  35 
  36     def _build_optimizer(
  37         self, optimizer_name: str, model, lr: float, weight_decay: float
  38     ) -> th.optim.Optimizer:
  39         """Build an optimizer if needed.
  40 
  41         Args:
  42             optimizer_name: A string indicating the optimizer name.
  43             lr: A float indicating the learning rate.
  44             weight_decay: Weight decay parameter of the optimizer
  45         Returns:
  46             A Torch Optimizer.
  47         """
  48         if self.optimizer is not None:
  49             return self.optimizer
  50 
  51         optimizer_name = optimizer_name.lower()
  52         if optimizer_name == "sgd":
  53             optim_args = dict()
  54             optim_args["lr"] = lr
  55             if weight_decay is not None:
  56                 optim_args["weight_decay"] = weight_decay
  57             self.optimizer = th.optim.SGD(model.parameters(), **optim_args)
  58         else:
  59             raise ValueError("Unknown optimizer: {}".format(optimizer_name))
  60         return self.optimizer
  61 
  62     def fit(self, dataset_key: str, **kwargs):
  63         """Fits a model on the local dataset as specified in the local TrainConfig object.
  64 
  65         Args:
  66             dataset_key: Identifier of the local dataset that shall be used for training.
  67             **kwargs: Unused.
  68 
  69         Returns:
  70             loss: Training loss on the last batch of training data.
  71         """
  72         if self.train_config is None:
  73             raise ValueError("TrainConfig not defined.")
  74 
  75         model = self.get_obj(self.train_config._model_id).obj
  76         loss_fn = self.get_obj(self.train_config._loss_fn_id).obj
  77 
  78         self._build_optimizer(
  79             self.train_config.optimizer,
  80             model,
  81             lr=self.train_config.lr,
  82             weight_decay=self.train_config.weight_decay,
  83         )
  84 
  85         return self._fit(model=model, dataset_key=dataset_key, loss_fn=loss_fn)
  86 
  87     def _create_data_loader(self, dataset_key: str, shuffle: bool = False):
  88         data_range = range(len(self.datasets[dataset_key]))
  89         if shuffle:
  90             sampler = RandomSampler(data_range)
  91         else:
  92             sampler = SequentialSampler(data_range)
  93         data_loader = th.utils.data.DataLoader(
  94             self.datasets[dataset_key],
  95             batch_size=self.train_config.batch_size,
  96             sampler=sampler,
  97             num_workers=0,
  98         )
  99         return data_loader
  100 
  101     def _fit(self, model, dataset_key, loss_fn):
  102         model.train()
  103         data_loader = self._create_data_loader(
  104             dataset_key=dataset_key, shuffle=self.train_config.shuffle
  105         )
  106 
  107         loss = None
  108         iteration_count = 0
  109         for (data, target) in data_loader:
  110             # Set gradients to zero
  111             self.optimizer.zero_grad()
  112 
  113             # Update model
  114             output = model(data)
  115             loss = loss_fn(target=target, pred=output)
  116             loss.backward()
  117             self.optimizer.step()
  118 
  119             # Update and check interation count
  120             iteration_count += 1
  121             if iteration_count >= self.train_config.max_nr_batches >= 0:
  122                 break
  123 
  124         return loss
  125
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/syft/federated/federated_client.py b/syft/federated/federated_client.py
  --- a/syft/federated/federated_client.py
  +++ b/syft/federated/federated_client.py
  @@ -72,6 +72,9 @@
           if self.train_config is None:
               raise ValueError("TrainConfig not defined.")
   
  +        if dataset_key not in self.datasets:
  +            raise ValueError("Dataset {} unknown.".format(dataset_key))
  +
           model = self.get_obj(self.train_config._model_id).obj
           loss_fn = self.get_obj(self.train_config._loss_fn_id).obj
   
  @@ -106,19 +109,21 @@
   
           loss = None
           iteration_count = 0
  -        for (data, target) in data_loader:
  -            # Set gradients to zero
  -            self.optimizer.zero_grad()
  -
  -            # Update model
  -            output = model(data)
  -            loss = loss_fn(target=target, pred=output)
  -            loss.backward()
  -            self.optimizer.step()
  -
  -            # Update and check interation count
  -            iteration_count += 1
  -            if iteration_count >= self.train_config.max_nr_batches >= 0:
  -                break
  +
  +        for _ in range(self.train_config.epochs):
  +            for (data, target) in data_loader:
  +                # Set gradients to zero
  +                self.optimizer.zero_grad()
  +
  +                # Update model
  +                output = model(data)
  +                loss = loss_fn(target=target, pred=output)
  +                loss.backward()
  +                self.optimizer.step()
  +
  +                # Update and check interation count
  +                iteration_count += 1
  +                if iteration_count >= self.train_config.max_nr_batches >= 0:
  +                    break
   
           return loss
  
verification_info: {"golden_diff": "diff --git a/syft/federated/federated_client.py b/syft/federated/federated_client.py\n--- a/syft/federated/federated_client.py\n+++ b/syft/federated/federated_client.py\n@@ -72,6 +72,9 @@\n         if self.train_config is None:\n             raise ValueError(\"TrainConfig not defined.\")\n \n+        if dataset_key not in self.datasets:\n+            raise ValueError(\"Dataset {} unknown.\".format(dataset_key))\n+\n         model = self.get_obj(self.train_config._model_id).obj\n         loss_fn = self.get_obj(self.train_config._loss_fn_id).obj\n \n@@ -106,19 +109,21 @@\n \n         loss = None\n         iteration_count = 0\n-        for (data, target) in data_loader:\n-            # Set gradients to zero\n-            self.optimizer.zero_grad()\n-\n-            # Update model\n-            output = model(data)\n-            loss = loss_fn(target=target, pred=output)\n-            loss.backward()\n-            self.optimizer.step()\n-\n-            # Update and check interation count\n-            iteration_count += 1\n-            if iteration_count >= self.train_config.max_nr_batches >= 0:\n-                break\n+\n+        for _ in range(self.train_config.epochs):\n+            for (data, target) in data_loader:\n+                # Set gradients to zero\n+                self.optimizer.zero_grad()\n+\n+                # Update model\n+                output = model(data)\n+                loss = loss_fn(target=target, pred=output)\n+                loss.backward()\n+                self.optimizer.step()\n+\n+                # Update and check interation count\n+                iteration_count += 1\n+                if iteration_count >= self.train_config.max_nr_batches >= 0:\n+                    break\n \n         return loss\n", "issue": "TrainConfig parameter \"epochs\"\n**TrainConfig parameter \"epochs\" doesn't have effect.**\r\nAfter changing the number of epochs=1 to epochs=100. The worker still do only 1 epoch.\r\n\r\n```\r\ntrain_config = sy.TrainConfig(\r\n        model=traced_model,\r\n        loss_fn=loss_fn,\r\n        batch_size=batch_size,\r\n        shuffle=True,\r\n        #max_nr_batches=max_nr_batches,\r\n        epochs=100,\r\n        lr=lr,\r\n    )\r\n```\n", "before_files": [{"content": "import torch as th\nfrom torch.utils.data import BatchSampler, RandomSampler, SequentialSampler\n\nfrom syft.generic import ObjectStorage\nfrom syft.federated.train_config import TrainConfig\n\n\nclass FederatedClient(ObjectStorage):\n    \"\"\"A Client able to execute federated learning in local datasets.\"\"\"\n\n    def __init__(self, datasets=None):\n        super().__init__()\n        self.datasets = datasets if datasets is not None else dict()\n        self.optimizer = None\n        self.train_config = None\n\n    def add_dataset(self, dataset, key: str):\n        self.datasets[key] = dataset\n\n    def remove_dataset(self, key: str):\n        if key in self.datasets:\n            del self.datasets[key]\n\n    def set_obj(self, obj: object):\n        \"\"\"Registers objects checking if which objects it should cache.\n\n        Args:\n            obj: An object to be registered.\n        \"\"\"\n        if isinstance(obj, TrainConfig):\n            self.train_config = obj\n            self.optimizer = None\n        else:\n            super().set_obj(obj)\n\n    def _build_optimizer(\n        self, optimizer_name: str, model, lr: float, weight_decay: float\n    ) -> th.optim.Optimizer:\n        \"\"\"Build an optimizer if needed.\n\n        Args:\n            optimizer_name: A string indicating the optimizer name.\n            lr: A float indicating the learning rate.\n            weight_decay: Weight decay parameter of the optimizer\n        Returns:\n            A Torch Optimizer.\n        \"\"\"\n        if self.optimizer is not None:\n            return self.optimizer\n\n        optimizer_name = optimizer_name.lower()\n        if optimizer_name == \"sgd\":\n            optim_args = dict()\n            optim_args[\"lr\"] = lr\n            if weight_decay is not None:\n                optim_args[\"weight_decay\"] = weight_decay\n            self.optimizer = th.optim.SGD(model.parameters(), **optim_args)\n        else:\n            raise ValueError(\"Unknown optimizer: {}\".format(optimizer_name))\n        return self.optimizer\n\n    def fit(self, dataset_key: str, **kwargs):\n        \"\"\"Fits a model on the local dataset as specified in the local TrainConfig object.\n\n        Args:\n            dataset_key: Identifier of the local dataset that shall be used for training.\n            **kwargs: Unused.\n\n        Returns:\n            loss: Training loss on the last batch of training data.\n        \"\"\"\n        if self.train_config is None:\n            raise ValueError(\"TrainConfig not defined.\")\n\n        model = self.get_obj(self.train_config._model_id).obj\n        loss_fn = self.get_obj(self.train_config._loss_fn_id).obj\n\n        self._build_optimizer(\n            self.train_config.optimizer,\n            model,\n            lr=self.train_config.lr,\n            weight_decay=self.train_config.weight_decay,\n        )\n\n        return self._fit(model=model, dataset_key=dataset_key, loss_fn=loss_fn)\n\n    def _create_data_loader(self, dataset_key: str, shuffle: bool = False):\n        data_range = range(len(self.datasets[dataset_key]))\n        if shuffle:\n            sampler = RandomSampler(data_range)\n        else:\n            sampler = SequentialSampler(data_range)\n        data_loader = th.utils.data.DataLoader(\n            self.datasets[dataset_key],\n            batch_size=self.train_config.batch_size,\n            sampler=sampler,\n            num_workers=0,\n        )\n        return data_loader\n\n    def _fit(self, model, dataset_key, loss_fn):\n        model.train()\n        data_loader = self._create_data_loader(\n            dataset_key=dataset_key, shuffle=self.train_config.shuffle\n        )\n\n        loss = None\n        iteration_count = 0\n        for (data, target) in data_loader:\n            # Set gradients to zero\n            self.optimizer.zero_grad()\n\n            # Update model\n            output = model(data)\n            loss = loss_fn(target=target, pred=output)\n            loss.backward()\n            self.optimizer.step()\n\n            # Update and check interation count\n            iteration_count += 1\n            if iteration_count >= self.train_config.max_nr_batches >= 0:\n                break\n\n        return loss\n", "path": "syft/federated/federated_client.py"}], "after_files": [{"content": "import torch as th\nfrom torch.utils.data import BatchSampler, RandomSampler, SequentialSampler\n\nfrom syft.generic import ObjectStorage\nfrom syft.federated.train_config import TrainConfig\n\n\nclass FederatedClient(ObjectStorage):\n    \"\"\"A Client able to execute federated learning in local datasets.\"\"\"\n\n    def __init__(self, datasets=None):\n        super().__init__()\n        self.datasets = datasets if datasets is not None else dict()\n        self.optimizer = None\n        self.train_config = None\n\n    def add_dataset(self, dataset, key: str):\n        self.datasets[key] = dataset\n\n    def remove_dataset(self, key: str):\n        if key in self.datasets:\n            del self.datasets[key]\n\n    def set_obj(self, obj: object):\n        \"\"\"Registers objects checking if which objects it should cache.\n\n        Args:\n            obj: An object to be registered.\n        \"\"\"\n        if isinstance(obj, TrainConfig):\n            self.train_config = obj\n            self.optimizer = None\n        else:\n            super().set_obj(obj)\n\n    def _build_optimizer(\n        self, optimizer_name: str, model, lr: float, weight_decay: float\n    ) -> th.optim.Optimizer:\n        \"\"\"Build an optimizer if needed.\n\n        Args:\n            optimizer_name: A string indicating the optimizer name.\n            lr: A float indicating the learning rate.\n            weight_decay: Weight decay parameter of the optimizer\n        Returns:\n            A Torch Optimizer.\n        \"\"\"\n        if self.optimizer is not None:\n            return self.optimizer\n\n        optimizer_name = optimizer_name.lower()\n        if optimizer_name == \"sgd\":\n            optim_args = dict()\n            optim_args[\"lr\"] = lr\n            if weight_decay is not None:\n                optim_args[\"weight_decay\"] = weight_decay\n            self.optimizer = th.optim.SGD(model.parameters(), **optim_args)\n        else:\n            raise ValueError(\"Unknown optimizer: {}\".format(optimizer_name))\n        return self.optimizer\n\n    def fit(self, dataset_key: str, **kwargs):\n        \"\"\"Fits a model on the local dataset as specified in the local TrainConfig object.\n\n        Args:\n            dataset_key: Identifier of the local dataset that shall be used for training.\n            **kwargs: Unused.\n\n        Returns:\n            loss: Training loss on the last batch of training data.\n        \"\"\"\n        if self.train_config is None:\n            raise ValueError(\"TrainConfig not defined.\")\n\n        if dataset_key not in self.datasets:\n            raise ValueError(\"Dataset {} unknown.\".format(dataset_key))\n\n        model = self.get_obj(self.train_config._model_id).obj\n        loss_fn = self.get_obj(self.train_config._loss_fn_id).obj\n\n        self._build_optimizer(\n            self.train_config.optimizer,\n            model,\n            lr=self.train_config.lr,\n            weight_decay=self.train_config.weight_decay,\n        )\n\n        return self._fit(model=model, dataset_key=dataset_key, loss_fn=loss_fn)\n\n    def _create_data_loader(self, dataset_key: str, shuffle: bool = False):\n        data_range = range(len(self.datasets[dataset_key]))\n        if shuffle:\n            sampler = RandomSampler(data_range)\n        else:\n            sampler = SequentialSampler(data_range)\n        data_loader = th.utils.data.DataLoader(\n            self.datasets[dataset_key],\n            batch_size=self.train_config.batch_size,\n            sampler=sampler,\n            num_workers=0,\n        )\n        return data_loader\n\n    def _fit(self, model, dataset_key, loss_fn):\n        model.train()\n        data_loader = self._create_data_loader(\n            dataset_key=dataset_key, shuffle=self.train_config.shuffle\n        )\n\n        loss = None\n        iteration_count = 0\n\n        for _ in range(self.train_config.epochs):\n            for (data, target) in data_loader:\n                # Set gradients to zero\n                self.optimizer.zero_grad()\n\n                # Update model\n                output = model(data)\n                loss = loss_fn(target=target, pred=output)\n                loss.backward()\n                self.optimizer.step()\n\n                # Update and check interation count\n                iteration_count += 1\n                if iteration_count >= self.train_config.max_nr_batches >= 0:\n                    break\n\n        return loss\n", "path": "syft/federated/federated_client.py"}]}