pylint fixes

Author: priyakasimbeg

algoperf/pytorch_utils.py

@@ -7,7 +7,7 @@
 import torch
 from torch import Tensor
 import torch.distributed as dist
-import torch.nn as nn
+from torch import nn
 import torch.nn.functional as F
 
 from algoperf import spec
@@ -100,8 +100,8 @@ def __init__(self):
     super().__init__()
     self._supports_custom_dropout = True
 
-  def forward(self, input: Tensor, p: float) -> Tensor:
-    return F.dropout2d(input, p, training=self.training)
+  def forward(self, x: Tensor, p: float) -> Tensor:
+    return F.dropout2d(x, p, training=self.training)
 
 
 class SequentialWithDropout(nn.Sequential):

algoperf/workloads/imagenet_vit/imagenet_jax/workload.py

@@ -1,6 +1,6 @@
 """ImageNet workload implemented in Jax."""
 
-from typing import Dict, Tuple
+from typing import Dict,  Optional, Tuple
 
 from flax import jax_utils
 from flax import linen as nn
@@ -54,10 +54,12 @@ def model_fn(
       mode: spec.ForwardPassMode,
       rng: spec.RandomState,
       update_batch_norm: bool,
+      use_running_average_bn: Optional[bool] = None,
       dropout_rate: float = models.DROPOUT_RATE
   ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
     del model_state
     del update_batch_norm
+    del use_running_average_bn
     train = mode == spec.ForwardPassMode.TRAIN
     logits = self._model.apply({'params': params},
                                augmented_and_preprocessed_input_batch['inputs'],