[DCU] fix batchgemm

transformer_engine/common/gemm/cublaslt_gemm.cu

@@ -731,20 +731,6 @@ void nvte_cublas_handle_init() { auto _ = cublasHandleManager::Instance().GetHan
 #endif
 
 #ifdef __HIP_PLATFORM_AMD__
-static inline int getIntEnv(const char *name, int defval, int minval)
-{
-  int val = defval;
-  const char* env = std::getenv(name);
-  if (env != nullptr && env[0] != '\0')
-  {
-     val = atoi(env);
-     if (val < minval)
-     {
-        val = minval;
-     }
-  }
-  return val;
-}
 
 void nvte_multi_stream_cublas_batchgemm(const NVTETensor *A, const NVTETensor *B, NVTETensor *D,
                                    const NVTETensor *bias, NVTETensor *pre_gelu_out,

transformer_engine/pytorch/module/__init__.py

@@ -5,7 +5,8 @@
 """Module level PyTorch APIs"""
 from .layernorm_linear import LayerNormLinear
 from .linear import Linear
-from .grouped_linear import GroupedLinear, BatchedLinear
+from .grouped_linear import GroupedLinear
+from .batched_linear import BatchedLinear
 from .layernorm_mlp import LayerNormMLP
 from .layernorm import LayerNorm
 from .rmsnorm import RMSNorm