Error when trying to use a custom dataset for conditional training

[Mathews-Tom]

Getting the following error when trying to use a custom dataset for conditional training.

TypeError: any() received an invalid combination of arguments - got (bool, dim=int), but expected one of:
 * (Tensor input, *, Tensor out)
 * (Tensor input, int dim, bool keepdim, *, Tensor out)
 * (Tensor input, name dim, bool keepdim, *, Tensor out)

Error occurs from this particular line

Using imagen-pytorch version 1.22.2.

PS: Somehow the line number mentioned in the error stack is different from the code in the repo