(self, a, axis=None)
| 2340 | ), f"Device discrepancy. First input is on {str(a_device)}, whereas second input is on {str(b_device)}" |
| 2341 | |
| 2342 | def squeeze(self, a, axis=None): |
| 2343 | if axis is None: |
| 2344 | return torch.squeeze(a) |
| 2345 | else: |
| 2346 | return torch.squeeze(a, dim=axis) |
| 2347 | |
| 2348 | def unsqueeze(self, a, axis): |
| 2349 | return torch.unsqueeze(a, dim=axis) |