code: def test_regression_2d(self, device):
      # create the parameters to estimate: the heatmap
      params = nn.Parameter(
          torch.rand(1, 1, self.height, self.width).to(device)
      )
  
      # generate base sample
      target = torch.zeros(1, 1, 2).to(device)
      target[..., 0] = self.width / 2
      target[..., 1] = self.height / 2
  
      # create the optimizer and pass the heatmap
      optimizer = optim.Adam([params], lr=self.lr)
  
      # loss criterion
      criterion = nn.MSELoss()
  
      # spatial soft-argmax2d module
      soft_argmax2d = kornia.geometry.SpatialSoftArgmax2d(
          normalized_coordinates=False
      )
  
      # NOTE: check where this comes from
      temperature = (self.height * self.width) ** (0.5)
  
      for _ in range(self.num_iterations):
          x = params
          sample = self.generate_sample(target).to(device)
          pred = soft_argmax2d(temperature * x)
          loss = criterion(pred, sample)
          logger.debug(f"Loss: {loss.item():.3f} Pred: {pred}")
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
      assert_close(pred[..., 0], target[..., 0], rtol=1e-2, atol=1e-2)
      assert_close(pred[..., 1], target[..., 1], rtol=1e-2, atol=1e-2)
  
file_path: kornia/test/integration/test_soft_argmax2d.py
imports: import logging
  import pytest
  import torch
  from torch import nn, optim
  import kornia
  from kornia.testing import assert_close
label_desc: Use image node_id 2 for calling the TestIntegrationSoftArgmax2d obj's underlying member method code with example usage: obj.test_regression_2d(device) without return types