def _to(x, device):
if hasattr(x, 'to'):
return x.to(device)
elif isinstance(x, torch.optim.Optimizer):
optim_sd = bd.apply(partial(_to, device=device))(x.state_dict())
x.load_state_dict(optim_sd)
return x
else:
return x
def test_nested_with_recurse(self):
a = {'a': [1, 2, (3, 4), {'b': 5, 'c': 6}], 'd': 7, 'e': [8, 9]}
b = {'a': [2, 3, (4, 5), {'b': 6, 'c': 7}], 'd': 8, 'e': [9, 10]}
assert bd.apply(add_one, recurse=True)(a) == b
def test_dict_with_recurse(self):
applier = bd.apply(add_one, recurse=True)
assert applier({'a': 1, 'b': 2}) == {'a': 2, 'b': 3}
def test_tuple_with_recurse(self):
assert bd.apply(add_one, recurse=True)((1, 2)) == (2, 3)
def test_list_with_recurse(self):
assert bd.apply(add_one, recurse=True)([1, 2]) == [2, 3]
def test_single_with_recurse(self):
assert bd.apply(add_one, recurse=True)(1) == 2
def test_dict(self):
assert bd.apply(add_one)({'a': 1, 'b': 2}) == {'a': 2, 'b': 3}
def test_tuple(self):
assert bd.apply(add_one)((1, 2)) == (2, 3)
def test_list(self):
assert bd.apply(add_one)([1, 2]) == [2, 3]
def test_single_value(self):
assert bd.apply(add_one)(1) == 2