A LightningModule has a self.log method for logging metrics. This takes the batch_size argument to help it know what your batch size is when you have a complex data structure. Our batch is a tuple of 7ish things, so by default lightning would see our batch size as 7. We get around this by explicitly passing batch_size=len(batch[0]). This calls the __len__ of the batch mol graphs which is set to the number of mol graphs.

MulticomponentMPNN inherits from MPNN and reuses a lot of its code. But batch[0] for the multicomponent case is a list of BatchMolGraphs. len(batch[0]) will be the number of components and not the true batch size. I fixed this by adding a get_batch_size method to both MulticomponentMPNN and MPNN. I also updated the type hints to reflect the fact that MulticomponentMPNN reuses code from MPNN.