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