(fn, tree, *rest)
| 56 | |
| 57 | @staticmethod |
| 58 | def tree_tensor_map(fn, tree, *rest): |
| 59 | return optree.tree_map( |
| 60 | fn, |
| 61 | tree, |
| 62 | *rest, |
| 63 | is_leaf=lambda x: isinstance(x, torch.Tensor), |
| 64 | none_is_leaf=False, |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | def to_device(obj, device): |
no outgoing calls
no test coverage detected