[bug] in write_image_imageio(), doesn't handle grayscale image write correctly

[ZJLi2013]

python3  samples/mlp_learning_an_image_pytorch.py  # with default albert.jpg & config.json

give errors:

ValueError: Can't write images with one color channel.

looks the write_image_imageio() under common.py missing handle grayscale image correctly. added as following:

def write_image_imageio(img_file, img, quality):
	img = (np.clip(img, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)   # 
	kwargs = {}
	if os.path.splitext(img_file)[1].lower() in [".jpg", ".jpeg"]:
		if img.ndim >= 3 and img.shape[2] > 3:
			img = img[:,:,:3]
		if img.ndim >= 3 and img.shape[2] == 1:  # grayscale 
			img = np.squeeze(img)
		kwargs["quality"] = quality
		kwargs["subsampling"] = 0
	imageio.imwrite(img_file, img, **kwargs)