Repository: kornia/kornia

Issue:
RandomPlanckianJitter incorrect values on gpu other than 0
### Describe the bug

`RandomPlanckianJitter` has a bug from `get_planckian_coeffs(mode)` that always puts `self.pl` on gpu 0.
`self.pl` is not registered in the module and does not get moved to the correct gpu.
Moving `self.pl `from gpu0 to another gpu directly results in incorrect values.

### Reproduction steps

```bash
augmentation = kornia.augmentation.RandomPlanckianJitter(p=1.0)
pl = augmentation.pl.clone()
augmentation(torch.rand(3, 3, 3).to("cuda:0"))
print((pl.cpu() == augmentation.pl.cpu()).all())
augmentation(torch.rand(3, 3, 3).to("cuda:1"))
print((pl.cpu() == augmentation.pl.cpu()).all())

-
tensor(True)
tensor(False)
```


### Expected behavior

Either self.pl should be registered so it moves when the module moves to the correct gpu or self.pl should remain on cpu and get moved in the call.

### Environment
This does not occur on all machines.
Happens on this machine with 3070s but not on one with A5000.


```shell
PyTorch version: 2.2.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Fedora Linux 37 (Thirty Seven) (x86_64)
GCC version: (GCC) 12.3.1 20230508 (Red Hat 12.3.1-1)
Clang version: Could not collect
CMake version: version 3.27.7
Libc version: glibc-2.36

Python version: 3.11.6 (main, Oct  3 2023, 00:00:00) [GCC 12.3.1 20230508 (Red Hat 12.3.1-1)] (64-bit runtime)
Python platform: Linux-6.5.12-100.fc37.x86_64-x86_64-with-glibc2.36
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3070
GPU 1: NVIDIA GeForce RTX 3070
GPU 2: NVIDIA GeForce RTX 3070
GPU 3: NVIDIA GeForce RTX 3070

Nvidia driver version: 545.29.06
cuDNN version: Probably one of the following:
/usr/lib64/libcudnn.so.8.9.7
/usr/lib64/libcudnn_adv_infer.so.8.9.7
/usr/lib64/libcudnn_adv_train.so.8.9.7
/usr/lib64/libcudnn_cnn_infer.so.8.9.7
/usr/lib64/libcudnn_cnn_train.so.8.9.7
/usr/lib64/libcudnn_ops_infer.so.8.9.7
/usr/lib64/libcudnn_ops_train.so.8.9.7
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      43 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             48
On-line CPU(s) list:                0-47
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen Threadripper 3960X 24-Core Processor
CPU family:                         23
Model:                              49
Thread(s) per core:                 2
Core(s) per socket:                 24
Socket(s):                          1
Stepping:                           0
Frequency boost:                    enabled
CPU(s) scaling MHz:                 81%
CPU max MHz:                        3800.0000
CPU min MHz:                        2200.0000
BogoMIPS:                           7599.98
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sev sev_es
Virtualization:                     AMD-V
L1d cache:                          768 KiB (24 instances)
L1i cache:                          768 KiB (24 instances)
L2 cache:                           12 MiB (24 instances)
L3 cache:                           128 MiB (8 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-47
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow: Mitigation; Safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.0
[pip3] triton==2.2.0
[conda] numpy                     1.24.2                   pypi_0    pypi
[conda] pytorch-triton            2.1.0+7d1a95b046          pypi_0    pypi
[conda] torch-tb-profiler         0.4.1                    pypi_0    pypi
[conda] torchvision               0.16.0.dev20230508+cu118          pypi_0    pypi
```


### Additional context

_No response_


Discussion:
@Modexus could you send a PR with the fix ?

Resolving patch:
diff --git a/kornia/augmentation/_2d/intensity/planckian_jitter.py b/kornia/augmentation/_2d/intensity/planckian_jitter.py
--- a/kornia/augmentation/_2d/intensity/planckian_jitter.py
+++ b/kornia/augmentation/_2d/intensity/planckian_jitter.py
@@ -3,11 +3,9 @@
 from kornia.augmentation import random_generator as rg
 from kornia.augmentation._2d.intensity.base import IntensityAugmentationBase2D
 from kornia.core import Tensor, stack, tensor
-from kornia.utils import get_cuda_or_mps_device_if_available
 
 
 def get_planckian_coeffs(mode: str) -> Tensor:
-    default_device = get_cuda_or_mps_device_if_available()
     if mode.lower() == "blackbody":
         coefs = tensor(
             [
@@ -36,8 +34,7 @@ def get_planckian_coeffs(mode: str) -> Tensor:
                 [0.3586, 0.4655, 0.6820],
                 [0.3544, 0.4654, 0.6878],
                 [0.3503, 0.4653, 0.6933],
-            ],
-            device=default_device,
+            ]
         )
 
     elif mode.upper() == "CIED":
@@ -66,8 +63,7 @@ def get_planckian_coeffs(mode: str) -> Tensor:
                 [0.3313, 0.4719, 0.6766],
                 [0.3263, 0.4719, 0.6826],
                 [0.3217, 0.4719, 0.6882],
-            ],
-            device=default_device,
+            ]
         )
     else:
         raise RuntimeError(f"Unexpected mode. Gotcha {mode}")
@@ -159,9 +155,8 @@ def __init__(
 
         _pl = get_planckian_coeffs(mode)
         if select_from is not None:
-            self.pl = _pl[select_from]
-        else:
-            self.pl = _pl
+            _pl = _pl[select_from]
+        self.register_buffer("pl", _pl)
 
         # the range of the sampling parameters
         _param_min: float = 0.0
@@ -174,8 +169,6 @@ def apply_transform(
     ) -> Tensor:
         list_idx = params["idx"].tolist()
 
-        self.pl = self.pl.to(device=input.device)
-
         coeffs = self.pl[list_idx]
 
         r_w = coeffs[:, 0][..., None, None]


Test patch: