(func, *args, unsqueeze=False)
| 39 | |
| 40 | |
| 41 | def wrap(func, *args, unsqueeze=False): |
| 42 | args = list(args) |
| 43 | for i, arg in enumerate(args): |
| 44 | if type(arg) == np.ndarray: |
| 45 | args[i] = torch.from_numpy(arg) |
| 46 | if unsqueeze: |
| 47 | args[i] = args[i].unsqueeze(0) |
| 48 | |
| 49 | result = func(*args) |
| 50 | |
| 51 | if isinstance(result, tuple): |
| 52 | result = list(result) |
| 53 | for i, res in enumerate(result): |
| 54 | if type(res) == torch.Tensor: |
| 55 | if unsqueeze: |
| 56 | res = res.squeeze(0) |
| 57 | result[i] = res.numpy() |
| 58 | return tuple(result) |
| 59 | elif type(result) == torch.Tensor: |
| 60 | if unsqueeze: |
| 61 | result = result.squeeze(0) |
| 62 | return result.numpy() |
| 63 | else: |
| 64 | return result |
| 65 | |
| 66 |
no outgoing calls
no test coverage detected