| 48 | |
| 49 | |
| 50 | class MapNestedTensorObjectImpl: |
| 51 | def __init__(self, tensor_map_fn): |
| 52 | self.tensor_map_fn = tensor_map_fn |
| 53 | |
| 54 | def __call__(self, object): |
| 55 | if isinstance(object, torch.Tensor): |
| 56 | return self.tensor_map_fn(object) |
| 57 | |
| 58 | elif isinstance(object, dict): |
| 59 | mapped_dict = {} |
| 60 | for key, value in object.items(): |
| 61 | mapped_dict[self(key)] = self(value) |
| 62 | return mapped_dict |
| 63 | |
| 64 | elif isinstance(object, (list, tuple)): |
| 65 | mapped_iter = [] |
| 66 | for iter in object: |
| 67 | mapped_iter.append(self(iter)) |
| 68 | return mapped_iter if not isinstance(object, tuple) else tuple(mapped_iter) |
| 69 | |
| 70 | else: |
| 71 | return object |
| 72 | |
| 73 | |
| 74 | def map_nested_tensor_object(object, tensor_map_fn): |
no outgoing calls
no test coverage detected