code: def main():
      device = torch.device(
          "cuda:0" if torch.cuda.is_available() else "cpu"
      )
  
      # load pretrain weights
      # download url: https://download.pytorch.org/models/resnet34-333f7ec4.pth
      model_weight_path = "./resnet34-pre.pth"
      assert os.path.exists(
          model_weight_path
      ), "file {} does not exist.".format(model_weight_path)
  
      # option1
      net = resnet34()
      net.load_state_dict(
          torch.load(model_weight_path, map_location=device)
      )
      # change fc layer structure
      in_channel = net.fc.in_features
      net.fc = nn.Linear(in_channel, 5)
  
imports: import os
  import torch
  import torch.nn
  from model import resnet34
label_desc: Use image node_id 1 for calling a global function with example usage: main() without return types