[BugFix] update TensorDictModule out_keys on assignment

[gtnv]

TLDR

Assigning TensorDictModule.out_keys changed the visible keys but left the keys used to write module outputs unchanged. Changing a module's output count therefore paired its outputs with stale keys and failed.

This PR updates both key lists on assignment while keeping select_out_keys limited to output filtering.

Fixes #1407

Testing

Direct reproduction

import torch
from torch import nn
from tensordict import TensorDict
from tensordict.nn import TensorDictModule

net = nn.Linear(3, 4)
module = TensorDictModule(net, in_keys=["in"], out_keys=["out"])

def split_output(module, args, output):
    return output, output.mean(dim=1)

net.register_forward_hook(split_output)
module.out_keys = ["out1", "out2"]

x = torch.randn(3, 3)
result = module(TensorDict({"in": x}, [3]))

assert module.out_keys == module.out_keys_source == ["out1", "out2"]
assert "out" not in result
assert result["out1"].shape == (3, 4)
assert result["out2"].shape == (3,)
torch.testing.assert_close(result["out2"], result["out1"].mean(dim=1))

out1, out2 = module(x)
torch.testing.assert_close(out2, out1.mean(dim=1))

module.select_out_keys("out1")
module.out_keys = ["out1", "out2"]
assert "out2" in module(TensorDict({"in": x}, [3]))

Test suites

pytest \
  test/nn/test_nn.py::TestTDModule::test_out_keys_setter \
  test/nn/test_nn.py::TestTDSequence::test_key_exclusion_constructor \
  test/nn/test_nn.py::TestSelectOutKeys -q
# 62 passed

pytest test/nn -q
# 611 passed, 

pytest test/compile/test_compile.py::TestCudaGraphs::test_tdmodule -q
# 2 passed

pre-commit run --files \
  tensordict/nn/common.py \
  tensordict/nn/sequence.py \
  test/nn/test_nn.py
# all hooks passed