Multi-temporal raster source visualizer fails when batch size is 1

[chrisjkuch]

raster-vision/rastervision_pytorch_learner/rastervision/pytorch_learner/dataset/visualizer/visualizer.py

Lines 122 to 127 in e4e10ad

	subfigs = fig.subfigures(nrows=batch_sz, ncols=1, hspace=0.0)
	subfig_axs = [
	subfig.subplots(
	nrows=T, ncols=params['subplot_args']['ncols'])
	for subfig in subfigs.flat
	]

This code fails when the batch size is 1 because the created subfigs doesn't have a flat property if there's only one row and one column. Not sure whether this should be fixed upstream in matplotlib...

Matplotlib version 3.7.1, rastervision version 0.21.2

import matplotlib.pyplot as plt
fig = plt.figure()
subfigs = fig.subfigures(nrows=2, ncols=1, hspace=0)
subfigs.flat
#> <numpy.flatiter object at 0x5575c63e68f0>

subfigs = fig.subfigures(nrows=1, ncols=1, hspace=0)
subfigs.flat
#> Traceback (most recent call last):
#>   File "<string>", line 1, in <module>
#> AttributeError: 'SubFigure' object has no attribute 'flat'

^{Created at 2023-10-11 17:12:31 CDT by reprexlite v0.5.0}

[AdeelH]

Looks like fig.subfigures() needs squeeze=False to act more consistently, just like plt.subplots(). Thanks for catching that!

[chrisjkuch]

Thanks for such a quick fix!