[FutureWarning] Resolving warnings caused by deprecated function usage.

[drivanov]

This PR provides a fix for the following warnings:

/usr/local/lib/python3.12/dist-packages/apex/contrib/cudnn_gbn/batch_norm.py:12: FutureWarning: 
`torch.cuda.amp.custom_fwd(args...)` is deprecated. Please use `torch.amp.custom_fwd(args..., device_type='cuda')` instead.
  @custom_fwd
/usr/local/lib/python3.12/dist-packages/apex/contrib/cudnn_gbn/batch_norm.py:24: FutureWarning: 
`torch.cuda.amp.custom_bwd(args...)` is deprecated. Please use `torch.amp.custom_bwd(args..., device_type='cuda')` instead.
  @custom_bwd