Fix (graph/equalize): cleanup and device management

[Giuseppe5]

No description provided.

[nickfraser]

Does it make sense to explicitly specify 'cpu' here in all cases? I assume this would have equalization go slower in cases where everything would fit on the GPU. What you had originally seems to make perfect sense - to keep everything in the model device memory.

Perhaps we should consider a context manager (*gasp!* Yes, I know...) which can be deployed, which overrides this original functionality to provide a "memory efficient" version.

In any case, I think such functionality is not vital and I'm happy for this to be merged in the meantime...

[Giuseppe5]

These operations are not compute heavy so I'm not expecting a considerable slowdown. I guess the issue is more that we're copying data to/from cpu several times which is not ideal.

The problem with the current implementation is that we do not account for the possibility of mismatched devices (e.g., GPU 0 and 1). We could try to be smarter and pick one device instead of always falling back on CPU, but again the main issue is data movement, not compute, so that wouldn't solve much.

Overall I'm happy with this solution since it's a good trade-off of complexity vs benefits.