code: def make_temporal_shift(
      net, n_segment, n_div=8, place="blockres", temporal_pool=False
  ):
      if temporal_pool:
          n_segment_list = [
              n_segment,
              n_segment // 2,
              n_segment // 2,
              n_segment // 2,
          ]
      else:
          n_segment_list = [n_segment] * 4
      assert n_segment_list[-1] > 0
  
      if isinstance(net, torchvision.models.ResNet):
          if place == "block":
  
              def make_block_temporal(stage, this_segment):
                  blocks = list(stage.children())
                  for i, b in enumerate(blocks):
                      blocks[i] = TemporalShift(
                          b, n_segment=this_segment, n_div=n_div
                      )
                  return nn.Sequential(*(blocks))
  
              net.layer1 = make_block_temporal(
                  net.layer1, n_segment_list[0]
              )
              net.layer2 = make_block_temporal(
                  net.layer2, n_segment_list[1]
              )
              net.layer3 = make_block_temporal(
                  net.layer3, n_segment_list[2]
              )
              net.layer4 = make_block_temporal(
                  net.layer4, n_segment_list[3]
              )
  
          elif "blockres" in place:
              n_round = 1
              if len(list(net.layer3.children())) >= 23:
                  n_round = 2
  
              def make_block_temporal(stage, this_segment):
                  blocks = list(stage.children())
                  for i, b in enumerate(blocks):
                      if i % n_round == 0:
                          blocks[i].conv1 = TemporalShift(
                              b.conv1,
                              n_segment=this_segment,
                              n_div=n_div,
                          )
                  return nn.Sequential(*blocks)
  
              net.layer1 = make_block_temporal(
                  net.layer1, n_segment_list[0]
              )
              net.layer2 = make_block_temporal(
                  net.layer2, n_segment_list[1]
              )
              net.layer3 = make_block_temporal(
                  net.layer3, n_segment_list[2]
              )
              net.layer4 = make_block_temporal(
                  net.layer4, n_segment_list[3]
              )
      else:
          raise NotImplementedError(place)
  
file_path: towhee/towhee/models/tsm/temporal_shift.py
imports: import torch
  from torch import nn
  import torch.nn.functional
  import torchvision
label_desc: Use image node_id 1 for calling a global function with example usage: make_temporal_shift(net, n_segment, n_div, place, temporal_pool) and returns: nn, nn