Fix missing move and bare new in pytree from_str_internal (pytorch#6803)

* Use std::variant to implement pytree Key

Pull Request resolved: pytorch#6701

Key was a struct that should've been a union; std::variant makes using a union much easier.
ghstack-source-id: 253128071
@exported-using-ghexport

Differential Revision: [D65575184](https://our.internmc.facebook.com/intern/diff/D65575184/)

* Fix missing move and bare new in pytree from_str_internal

Just a couple minor fixes.

Differential Revision: [D65576543](https://our.internmc.facebook.com/intern/diff/D65576543/)

ghstack-source-id: 253128072
Pull Request resolved: pytorch#6783

---------

Co-authored-by: Scott Wolchok <[email protected]>

Author: pytorchbot

extension/pytree/pytree.h

@@ -139,6 +139,9 @@ struct ContainerHandle {
 
   /*implicit*/ ContainerHandle(container_type* c) : handle(c) {}
 
+  /*implicit*/ ContainerHandle(std::unique_ptr<container_type> c)
+      : handle(std::move(c)) {}
+
   void set_leaf(leaf_type* leaf) {
     pytree_assert(handle->kind == Kind::Leaf);
     handle->leaf = leaf;
@@ -486,10 +489,10 @@ TreeSpec<Aux> from_str_internal(
       read_idx++;
       auto layout = read_node_layout(spec, read_idx);
       const auto size = layout.size();
-      auto c = new TreeSpecContainer<Aux>(kind, size);
+      auto c = std::make_unique<TreeSpecContainer<Aux>>(kind, size);
 
       if (Kind::Custom == kind) {
-        c->custom_type = custom_type;
+        c->custom_type = std::move(custom_type);
       }
 
       size_t child_idx = 0;
@@ -509,14 +512,14 @@ TreeSpec<Aux> from_str_internal(
         read_idx++;
       }
       c->leaves_num = leaves_offset;
-      return c;
+      return TreeSpec<Aux>(std::move(c));
     }
 
     case Config::kDict: {
       read_idx++;
       auto layout = read_node_layout(spec, read_idx);
       const auto size = layout.size();
-      auto c = new TreeSpecContainer<Aux>(Kind::Dict, size);
+      auto c = std::make_unique<TreeSpecContainer<Aux>>(Kind::Dict, size);
 
       size_t child_idx = 0;
       size_t leaves_offset = 0;
@@ -549,7 +552,7 @@ TreeSpec<Aux> from_str_internal(
         read_idx++;
       }
       c->leaves_num = leaves_offset;
-      return c;
+      return TreeSpec<Aux>(std::move(c));
     }
 
     case Config::kLeaf: