[export][ez] Update tag_ for union setters. (pytorch#146912)

Summary: ez fix to set tag for union type fields.

Test Plan: CI

Differential Revision: D69467715

Pull Request resolved: pytorch#146912
Approved by: https://github.com/yiming0416

Author: zhxchen17

torch/_export/serde/schema_check.py

@@ -268,6 +268,7 @@ def accessor(name, ty, idx):
 
   void set_{name}({ty} def) {{
     variant_.emplace<{idx + 1}>(std::move(def));
+    tag_ = Tag::{name.upper()};
   }}
 """