(tensor)
| 45 | |
| 46 | |
| 47 | def _cast_if_autocast_enabled(tensor): |
| 48 | if torch.is_autocast_enabled(): |
| 49 | if tensor.device.type == "cuda": |
| 50 | dtype = torch.get_autocast_gpu_dtype() |
| 51 | elif tensor.device.type == "cpu": |
| 52 | dtype = torch.get_autocast_cpu_dtype() |
| 53 | else: |
| 54 | raise NotImplementedError() |
| 55 | return tensor.to(dtype=dtype) |
| 56 | return tensor |
| 57 | |
| 58 | |
| 59 | class LPLayerNorm(torch.nn.LayerNorm): |