Fixed mean_cache data extraction for dKG

[yyexela]

Issue: dKG in BoTorch no longer works since I last ran it a few months ago. It appears to be from the mean_cache going from size [num_fantasies, 1, 1, num_data] to [num_fantasies, num_fantasies, num_fantasies, num_data]. What this change does is extract the relevant data of mean_cache while keeping the expected shape for the rest of the code (ie. [num_fantasies, 1, num_data]).

Not sure about a unit test for this since it depends on having BoTorch's qKnowledgeGradient class, among other dependencies.

See meta-pytorch/botorch#2137

[yyexela]

Any idea why the unit test is failing when there is no code change?

[yyexela]

Noticed main fails the same test so I'm assuming my code works and it's another problem.

EDIT: Unit tests aren't passing but neither are the ones in main, so I'll assume this change is what's needed for now.

[CedricTravelletti]

Any news on this PR? I confirm that there is interest in gradient-enabled GPs/BO for several application cases.

[nrummel]

I am also hoping to have this change in the main branch, as I would like to use derivative information for Bayesian Optimization through BoTorch.

[yyexela]

@gpleiss any way we can get this through? It's a simple fix that would allow people to run this botorch notebook, there's clearly an interest in the community to do work on dGPs (meta-pytorch/botorch#2137).