code: def test_non_param_state_key(self) -> None:
      with self.assertRaisesRegex(
          ValueError, "All state keys must be params."
      ):
          param_1_t = torch.tensor([1.0, 2.0])
          param_1 = Variable(param_1_t)
          KeyedOptimizer(
              {"param_1": param_1},
              {param_1: 1.0, "non_param_state_key": 2.0},
              [{"params": [param_1], "param_group_val_0": 3.0}],
          )
  
imports: import io
  import os
  import unittest
  from typing import Any, Dict, List
  import torch
  import torch.distributed
  from torch.autograd import Variable
  from torch.distributed._shard import sharded_tensor, sharding_spec
  from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizer, KeyedOptimizerWrapper, OptimizerWrapper
  from torchrec.test_utils import get_free_port
label_desc: Use image node_id 3 for calling the TestKeyedOptimizer obj's underlying member method code with example usage: obj.test_non_param_state_key() without return types