[fix] Reverting a change which slipped in #188 (#198)

* reverting a change which slipped in #188

benchmarks/pipe.py

@@ -685,8 +685,7 @@ parser.set_defaults(pipelined_backward=True)
 
 if __name__ == "__main__":
     args = parser.parse_args()
-    bench_multi_process(args, all_at_once=True)
-
+    # bench_multi_process(args, all_at_once=True)
     if args.no_mpi or "OMPI_COMM_WORLD_RANK" not in os.environ:
         print(f"Running benchmark with args: {args}")
         bench_single_process(args)

fairscale/nn/pipe/pipe.py

@@ -736,7 +736,10 @@ class Pipe(Module):
 
                     from .phony import get_phony
 
-                    phony = get_phony(torch.device(torch.cuda.current_device() if torch.cuda.is_available() else "cpu"), requires_grad=True)
+                    phony = get_phony(
+                        torch.device(torch.cuda.current_device() if torch.cuda.is_available() else "cpu"),
+                        requires_grad=True,
+                    )
                     output = PipelinedBackwardPass.apply(output, batches, phony, True)  # self.retain_graph)
                 else:
                     output = microbatch.gather(batches)