Random crashes: 'NoneType' object has no attribute 'ndim'

[sampbarrow]

Training is intermittently crashing on some sample generations. Error:

[rank3]: Traceback (most recent call last):
[rank3]: File "/root/run/kohya/flux_train.py", line 905, in
[rank3]: train(args)
[rank3]: File "/root/run/kohya/flux_train.py", line 760, in train
[rank3]: flux_train_utils.sample_images(
[rank3]: File "/root/run/kohya/library/flux_train_utils.py", line 113, in sample_images
[rank3]: sample_image_inference(
[rank3]: File "/root/run/kohya/library/flux_train_utils.py", line 229, in sample_image_inference
[rank3]: x = denoise(flux, noise, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=scale, t5_attn_mask=t5_attn_mask)
[rank3]: File "/root/run/kohya/library/flux_train_utils.py", line 314, in denoise
[rank3]: pred = model(
[rank3]: File "/root/miniconda/envs/kohya/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
[rank3]: return self._call_impl(*args, **kwargs)
[rank3]: File "/root/miniconda/envs/kohya/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
[rank3]: return forward_call(*args, **kwargs)
[rank3]: File "/root/miniconda/envs/kohya/lib/python3.10/site-packages/accelerate/utils/operations.py", line 819, in forward
[rank3]: return model_forward(*args, **kwargs)
[rank3]: File "/root/miniconda/envs/kohya/lib/python3.10/site-packages/accelerate/utils/operations.py", line 807, in call
[rank3]: return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank3]: File "/root/miniconda/envs/kohya/lib/python3.10/site-packages/torch/amp/autocast_mode.py", line 43, in decorate_autocast
[rank3]: return func(*args, **kwargs)
[rank3]: File "/root/run/kohya/library/flux_models.py", line 1004, in forward
[rank3]: if img.ndim != 3 or txt.ndim != 3:
[rank3]: AttributeError: 'NoneType' object has no attribute 'ndim'