How to get pytorch concat function to work by index


I have two models that take in images from a house, one being the house exterior and one being the roof. Each observation has multiple photos and not every observation has the same amount of photos. Is it possible to concat the two models by house ID?

If so, how would I do that. My current custom dataclass object returns the image, target label and ID.