(self)
| 60 | assert (layer.v == layer.rest * torch.ones(n)).all() |
| 61 | |
| 62 | def test_transfer(self): |
| 63 | if not torch.cuda.is_available(): |
| 64 | return |
| 65 | |
| 66 | for nodes in Nodes.__subclasses__(): |
| 67 | layer = nodes(10) |
| 68 | |
| 69 | layer.to(torch.device("cuda:0")) |
| 70 | |
| 71 | layer_tensors = [ |
| 72 | k for k, v in layer.state_dict().items() if isinstance(v, torch.Tensor) |
| 73 | ] |
| 74 | |
| 75 | tensor_devs = [getattr(layer, k).device for k in layer_tensors] |
| 76 | |
| 77 | print("State dict in {} : {}".format(nodes, layer.state_dict().keys())) |
| 78 | print("__dict__ in {} : {}".format(nodes, layer.__dict__.keys())) |
| 79 | print("Tensors in {} : {}".format(nodes, layer_tensors)) |
| 80 | print("Tensor devices {}".format(list(zip(layer_tensors, tensor_devs)))) |
| 81 | |
| 82 | for d in tensor_devs: |
| 83 | print(d, d == torch.device("cuda:0")) |
| 84 | assert d == torch.device("cuda:0") |
| 85 | |
| 86 | print("Reset layer") |
| 87 | layer.reset_state_variables() |
| 88 | layer_tensors = [ |
| 89 | k for k, v in layer.state_dict().items() if isinstance(v, torch.Tensor) |
| 90 | ] |
| 91 | |
| 92 | tensor_devs = [getattr(layer, k).device for k in layer_tensors] |
| 93 | |
| 94 | for d in tensor_devs: |
| 95 | print(d, d == torch.device("cuda:0")) |
| 96 | assert d == torch.device("cuda:0") |
| 97 | |
| 98 | |
| 99 | if __name__ == "__main__": |
no test coverage detected