in train.py [0:0]
def evaluate(H, ema_vae, data_valid, preprocess_fn):
stats_valid = []
valid_sampler = DistributedSampler(data_valid, num_replicas=H.mpi_size, rank=H.rank)
for x in DataLoader(data_valid, batch_size=H.n_batch, drop_last=True, pin_memory=True, sampler=valid_sampler):
data_input, target = preprocess_fn(x)
stats_valid.append(eval_step(data_input, target, ema_vae))
vals = [a['elbo'] for a in stats_valid]
finites = np.array(vals)[np.isfinite(vals)]
stats = dict(n_batches=len(vals), filtered_elbo=np.mean(finites), **{k: np.mean([a[k] for a in stats_valid]) for k in stats_valid[-1]})
return stats