Update main.py

atong01 · web-flow · commit b14f57b21e5c · 2020-10-15T08:15:31.000-04:00
Remove premature return statement in plot_output, now should attempt to plot both the trajectory and trajectory_density videos on each training run.
diff --git a/TrajectoryNet/main.py b/TrajectoryNet/main.py
@@ -389,7 +389,6 @@ def plot_output(device, args, model):
         end_times=args.int_tps,
         ntimes=100,
     )
-    return
     save_trajectory(
         args.data.base_density(),
         args.data.base_sample(),

Original file line number	Diff line number	Diff line change
`@@ -389,7 +389,6 @@ def plot_output(device, args, model):`
`389`	`389`	`end_times=args.int_tps,`
`390`	`390`	`ntimes=100,`
`391`	`391`	`)`
`392`		`- return`
`393`	`392`	`save_trajectory(`
`394`	`393`	`args.data.base_density(),`
`395`	`394`	`args.data.base_sample(),`