docstr: Do randomize to node_info.output_origin and store to node_info.output_randomize
docstr_tok: ["Do","randomize","to","node_info.output_origin","and","store","to","node_info.output_randomize"]
code: def direct_update_preprocess(
      self, model_speedup: "ModelSpeedup", node: Node
  ):
      """
      Do randomize to node_info.output_origin and store to node_info.output_randomize
      """
      node_info = model_speedup.node_infos[node]
      batch_dim, batch_size = (
          model_speedup.batch_dim,
          model_speedup.batch_size,
      )
      node_info.output_randomize = tree_map_zip(
          lambda t: randomize_if_tensor(t, batch_dim, batch_size),
          node_info.output_origin,
      )
  
imports: from __future__ import annotations
  from typing import TYPE_CHECKING, Tuple
  import operator
  import torch
  from torch.nn import functional
  from torch.fx.node import Node
  from torch.utils._pytree import tree_flatten, tree_unflatten
  from .utils import randomize_tensor_inplace, randomize_if_tensor, tree_map_zip, torch_float_dtype, poss_deepcopy
label_desc: Use image node_id 2 for calling the DefaultMaskUpdater obj's underlying member method code with example usage: obj.direct_update_preprocess(model_speedup, node) without return types