def METHOD_NAME(self):
    dtypes = [torch.float16, torch.float32, torch.bfloat16]
    for dtype in dtypes:
        weight = torch.randn(2, 3, dtype=dtype)
        with TemporaryDirectory() as tmp_dir:
            index = offload_weight(weight, "weight", tmp_dir, {})
            weight_file = os.path.join(tmp_dir, "weight.dat")
            self.assertTrue(os.path.isfile(weight_file))
            self.assertDictEqual(index, {"weight": {"shape": [2, 3], "dtype": str(dtype).split(".")[1]}})
            new_weight = load_offloaded_weight(weight_file, index["weight"])
            self.assertTrue(torch.equal(weight, new_weight))