torch.nn.modules.module.ModuleAttributeError: 'Labelator' object has no attribute 'decode'

[songweige]

trafficstars

I was trying to reproduce the label-conditioned model on ImageNet and ran into the error in the title. I think this line in the code, https://github.com/CompVis/taming-transformers/blob/9d17ea64b820f7633ea6b8823e1f78729447cb57/taming/models/cond_transformer.py#L242 should be replaced as:

if self.cond_stage_key not in ["image", "class_label"]:

Please let me know if that is the correct way to do or I missed something. Thank you!