Fix device handling in tests

[dtch1997]

Modifies ModelPatcher so that target activations are moved onto the appropriate device before performing operator.

closes #51

[dtch1997]

Setting the device here allows us to skip setting the activation devices.

E.g.:

activations = {
    1: torch.randn(1, 512), # can be on cpu
}
... 
model_patcher.patch_activations(
    activations, layer_type=layer_type, operator="piecewise_addition"
)

[dtch1997]

As of 16426e4, all the tests pass for me locally (with GPU enabled):

$ python -m pytest tests/util/test_model_patcher.py 
...
12 passed, 9 warnings in 15.54s

[chanind]

makes sense, LGTM!

[dtch1997]

@chanind after merging in your latest changes from main, one of the tests isn't working any more.

Looking at the error logs, it seems to be in this code chunk:

def _create_additive_hook(
    target_activation: torch.Tensor, operator: PatchOperator
) -> Any:
    """Create a hook function that adds the given target_activation to the model output"""

    def hook_fn(_m: Any, _inputs: Any, outputs: Any) -> Any:
        original_tensor = untuple_tensor(outputs)
        act = target_activation.to(original_tensor.device) # This line raises an error
        original_tensor[None] = operator(original_tensor, act)
        return outputs

    return hook_fn

Looking at the CI logs here, this line fails:

act = target_activation.to(original_tensor.device)

with this error:

AttributeError: 'numpy.ndarray' object has no attribute 'to'

I.e. the target_activation seems to be a numpy array. Is this behaviour intended? The type hint indicates that target_activation should be a torch.Tensor

[chanind]

Good catch! Fixed in #55