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 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.