def test_dp_output_reduce():
mixin = TrainerLoggingMixin()
# test identity when we have a single gpu
out = torch.rand(3, 1)
assert mixin.reduce_distributed_output(out, num_gpus=1) is out
# average when we have multiples
assert mixin.reduce_distributed_output(out, num_gpus=2) == out.mean()
# when we have a dict of vals
out = {'a': out, 'b': {'c': out}}
reduced = mixin.reduce_distributed_output(out, num_gpus=3)
assert reduced['a'] == out['a']
assert reduced['b']['c'] == out['b']['c']