Is it possible to save arbitrary data structures in test_epoch_end that are accessible from trainer.test(...)? I want to run my model over a bunch of data which produces nested tensors that I want to get a handle on. I see there is result.write however this only deals with flat tensors which is too restrictive for my use case