content: import os
  
  import hydra
  import hydra.experimental
  import numpy as np
  import pytest
  import torch
  
  pytest_plugins = ["helpers_namespace"]
  
  hydra.experimental.initialize(
      os.path.join(os.path.dirname(__file__), "../pointnet2/config")
  )
  
  
  @pytest.helpers.register
  def build_cfg(overrides=[]):
      return hydra.experimental.compose("config.yaml", overrides)
  
  
  @pytest.helpers.register
  def get_model(overrides=[]):
      cfg = build_cfg(overrides)
      return hydra.utils.instantiate(cfg.task_model, cfg)
  
  
  def _test_loop(model, inputs, labels):
      optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
  
      prev_loss = 1e10
      for _ in range(5):
          optimizer.zero_grad()
          res = model.training_step((inputs, labels), None)
          loss = res["loss"]
          loss.backward()
          optimizer.step()
  
          assert loss.item() < prev_loss + 1.0, "Loss spiked upwards"
  
          prev_loss = loss.item()
  
  
  @pytest.helpers.register
  def cls_test(model):
      B, N = 4, 2048
      inputs = torch.randn(B, N, 6).cuda()
      labels = torch.from_numpy(np.random.randint(0, 3, size=B)).cuda()
      model.cuda()
  
      _test_loop(model, inputs, labels)
  
  
  @pytest.helpers.register
  def semseg_test(model):
      B, N = 4, 2048
      inputs = torch.randn(B, N, 9).cuda()
      labels = torch.from_numpy(np.random.randint(0, 3, size=B * N)).view(B, N).cuda()
      model.cuda()
  
      _test_loop(model, inputs, labels)
  