hunk.code_window:             update_summary(
                   epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
                   write_header=best_metric is None, log_wandb=args.log_wandb)
               if saver is not None:
                   # save proper checkpoint with eval metric
  
hunk.labels: keep keep keep replace keep keep keep
hunk.after_edit:                 write_header=best_metric is None, log_wandb=args.log_wandb and has_wandb)
  
hunk.file_path: train.py
hunk.type: replace
file: import pytest
  import torch
  import torch.nn as nn
  import platform
  import os
  
  from timm.models.layers import create_act_layer, get_act_layer, set_layer_config
  
  
  class MLP(nn.Module):
      def __init__(self, act_layer="relu"):
          super(MLP, self).__init__()
          self.fc1 = nn.Linear(1000, 100)
          self.act = create_act_layer(act_layer, inplace=True)
          self.fc2 = nn.Linear(100, 10)
  
      def forward(self, x):
          x = self.fc1(x)
          x = self.act(x)
          x = self.fc2(x)
          return x
  
  
  def _run_act_layer_grad(act_type):
      x = torch.rand(10, 1000) * 10
      m = MLP(act_layer=act_type)
  
      def _run(x, act_layer=''):
          if act_layer:
              # replace act layer if set
              m.act = create_act_layer(act_layer, inplace=True)
          out = m(x)
          l = (out - 0).pow(2).sum()
          return l
  
      out_me = _run(x)
  
      with set_layer_config(scriptable=True):
          out_jit = _run(x, act_type)
  
      assert torch.isclose(out_jit, out_me)
  
      with set_layer_config(no_jit=True):
          out_basic = _run(x, act_type)
  
      assert torch.isclose(out_basic, out_jit)
  
  
  def test_swish_grad():
      for _ in range(100):
          _run_act_layer_grad('swish')
  
  
  def test_mish_grad():
      for _ in range(100):
          _run_act_layer_grad('mish')
  
  
  def test_hard_sigmoid_grad():
      for _ in range(100):
          _run_act_layer_grad('hard_sigmoid')
  
  
  def test_hard_swish_grad():
      for _ in range(100):
          _run_act_layer_grad('hard_swish')
  
  
  def test_hard_mish_grad():
      for _ in range(100):
          _run_act_layer_grad('hard_mish')
  