losses = losses.reshape(-1, args.maze_height, args.maze_width)
input = input.reshape(-1, args.maze_height, args.maze_width)
maze.save_image(
losses = losses.reshape(-1, args.maze_height, args.maze_width)
input = input.reshape(-1, args.maze_height, args.maze_width)
maze.save_image(