Repository: huggingface/pytorch-image-models
Issue: [BUG] tf_efficientnet_cc_b0_4e model doesn't work well with channels-last

**Describe the bug**
tf_efficientnet_cc_b0_4e model doesn't work well with channels-last

**To Reproduce**
Steps to reproduce the behavior:
```bash
python benchmark.py --model tf_efficientnet_cc_b0_4e -b 128 --amp --channels-last
```

Error message

```
WARNING: Overriding precision to 'amp' since --amp flag set.
Benchmarking in amp precision. NHWC layout. torchscript disabled
Model tf_efficientnet_cc_b0_4e created, param count: 13314116
Running inference benchmark on tf_efficientnet_cc_b0_4e for 40 steps w/ input size (3, 224, 224) and batch size 128.
WARNING: "view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead." while running benchmark. Reducing batch size to 128 for retry.
Model tf_efficientnet_cc_b0_4e created, param count: 13314116
Running inference benchmark on tf_efficientnet_cc_b0_4e for 40 steps w/ input size (3, 224, 224) and batch size 64.
WARNING: "view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead." while running benchmark. Reducing batch size to 64 for retry.
Model tf_efficientnet_cc_b0_4e created, param count: 13314116
Running inference benchmark on tf_efficientnet_cc_b0_4e for 40 steps w/ input size (3, 224, 224) and batch size 32.
WARNING: "view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead." while running benchmark. Reducing batch size to 32 for retry.
Model tf_efficientnet_cc_b0_4e created, param count: 13314116
Running inference benchmark on tf_efficientnet_cc_b0_4e for 40 steps w/ input size (3, 224, 224) and batch size 31.
WARNING: "view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead." while running benchmark. Reducing batch size to 31 for retry.

...
(and it keeps decreasing the batch size until 1, and gives a throughput, for example)
Running inference benchmark on tf_efficientnet_cc_b0_4e for 40 steps w/ input size (3, 224, 224) and batch size 1.
Infer [8/40]. 53.05 samples/sec. 18.850 ms/step.
Infer [16/40]. 52.62 samples/sec. 19.006 ms/step.
Infer [24/40]. 52.45 samples/sec. 19.067 ms/step.
Infer [32/40]. 52.45 samples/sec. 19.067 ms/step.
Infer [40/40]. 52.46 samples/sec. 19.064 ms/step.
...
(same issue for training mode)
```

**Expected behavior**
It should work.

**Screenshots**
N/A

**Desktop (please complete the following information):**
 - OS: linux
 - This repository version: 7c67d6aca992f039eece0af5f7c29a43d48c00e4
 - PyTorch version w/ CUDA/cuDNN: pytorch/pytorch@9b203f6, cuda 11.6, cudnn 8.3.2
 - GPU: V100, A100

**Additional context**
N/A

cc @ptrblck


Changed files:
timm/models/layers/cond_conv2d.py, timm/utils/jit.py

diff --git a/timm/models/layers/cond_conv2d.py b/timm/models/layers/cond_conv2d.py
index 8b4bbca..43654c5 100644
--- a/timm/models/layers/cond_conv2d.py
+++ b/timm/models/layers/cond_conv2d.py
@@ -91,7 +91,8 @@ class CondConv2d(nn.Module):
             bias = torch.matmul(routing_weights, self.bias)
             bias = bias.view(B * self.out_channels)
         # move batch elements with channels so each batch element can be efficiently convolved with separate kernel
-        x = x.view(1, B * C, H, W)
+        # reshape instead of view to work with channels_last input
+        x = x.reshape(1, B * C, H, W)
         if self.dynamic_padding:
             out = conv2d_same(
                 x, weight, bias, stride=self.stride, padding=self.padding,
diff --git a/timm/utils/jit.py b/timm/utils/jit.py
index 6039823..8ebfdbf 100644
--- a/timm/utils/jit.py
+++ b/timm/utils/jit.py
@@ -34,9 +34,9 @@ def set_jit_fuser(fuser):
         torch._C._jit_override_can_fuse_on_gpu(True)
         torch._C._jit_set_texpr_fuser_enabled(False)
     elif fuser == "nvfuser" or fuser == "nvf":
-        os.environ['PYTORCH_CUDA_FUSER_DISABLE_FALLBACK'] = '1'
-        os.environ['PYTORCH_CUDA_FUSER_DISABLE_FMA'] = '1'
-        os.environ['PYTORCH_CUDA_FUSER_JIT_OPT_LEVEL'] = '0'
+        os.environ['PYTORCH_NVFUSER_DISABLE_FALLBACK'] = '1'
+        os.environ['PYTORCH_NVFUSER_DISABLE_FMA'] = '1'
+        os.environ['PYTORCH_NVFUSER_JIT_OPT_LEVEL'] = '0'
         torch._C._jit_set_texpr_fuser_enabled(False)
         torch._C._jit_set_profiling_executor(True)
         torch._C._jit_set_profiling_mode(True)