# Buggy code
def train_epoch(self, *pargs, **kwargs):
    # print(self.model)
    def benchmark():
        self.optimizer.zero_grad()
        output = self.model(self.data)
        loss = F.cross_entropy(output, self.target)
        loss.backward()
        self.optimizer.step()

    # print("Running warmup...")
    if self.global_step == 0:
        timeit.timeit(benchmark, number=args.num_warmup_batches)
        self.global_step += 1
    # print("Running benchmark...")
    time = timeit.timeit(benchmark, number=args.num_batches_per_iter)
    img_sec = args.batch_size * args.num_batches_per_iter / time
    return {"img_sec": img_sec}


# Traceback (FileNotFoundError)
$ python dcgan.py --num-workers 44
2020-04-13 16:32:44,305 INFO resource_spec.py:204 -- Starting Ray with 120.21 GiB memory available for workers and up to 55.52 GiB for objects. You can adjust these settings with ray.init(me
mory=<bytes>, object_store_memory=<bytes>).
2020-04-13 16:32:44,732 INFO services.py:1146 -- View the Ray dashboard at localhost:8265
Traceback (most recent call last):
File "dcgan.py", line 283, in <module>
trainer = train_example(
File "dcgan.py", line 236, in train_example
trainer = TorchTrainer(
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/torch_trainer.py", line 233, in __init__
self._start_workers(self.max_replicas)
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/torch_trainer.py", line 320, in _start_workers
self.local_worker.setup(address, 0, num_workers)
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/distributed_torch_runner.py", line 46, in setup
self._setup_training()
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/distributed_torch_runner.py", line 92, in _setup_training
self.training_operator = self.training_operator_cls(
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/training_operator.py", line 96, in __init__
self.setup(config)
File "dcgan.py", line 136, in setup
torch.load(config["classification_model_path"]))
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 525, in load
with _open_file_like(f, 'rb') as opened_file:
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 212, in _open_file_like
return _open_file(name_or_buffer, mode)
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 193, in __init__
super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/examples/mnist_cnn.pt'
2020-04-13 16:32:49,606 ERROR worker.py:1011 -- Possible unhandled error from worker: ray::DistributedTorchRunner.setup() (pid=67029, ip=10.125.21.189)
File "python/ray/_raylet.pyx", line 452, in ray._raylet.execute_task
File "python/ray/_raylet.pyx", line 407, in ray._raylet.execute_task.function_executor
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/distributed_torch_runner.py", line 46, in setup
self._setup_training()
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/distributed_torch_runner.py", line 92, in _setup_training
self.training_operator = self.training_operator_cls(
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/training_operator.py", line 96, in __init__
self.setup(config)
File "dcgan.py", line 136, in setup
torch.load(config["classification_model_path"]))
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 525, in load
with _open_file_like(f, 'rb') as opened_file:
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 212, in _open_file_like
return _open_file(name_or_buffer, mode)
File "/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/torch/serialization.py", line 193, in __init__
super(_open_file, self).__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/rpovelik/installed/miniconda3/envs/ray/lib/python3.8/site-packages/ray/util/sgd/torch/examples/mnist_cnn.pt'

# Fixed code
def train_epoch(self, *pargs, **kwargs):
    def benchmark():
        self.optimizer.zero_grad()
        output = self.model(self.data)
        loss = F.cross_entropy(output, self.target)
        loss.backward()
        self.optimizer.step()

    print("Running warmup...")
    if self.global_step == 0:
        timeit.timeit(benchmark, number=args.num_warmup_batches)
        self.global_step += 1
    print("Running benchmark...")
    time = timeit.timeit(benchmark, number=args.num_batches_per_iter)
    img_sec = args.batch_size * args.num_batches_per_iter / time
    return {"img_sec": img_sec}