import torch
import torch.nn as nn

class ModelNew(nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A, B):
        # Optimized matrix multiplication using PyTorch operations
        # This implementation uses optimized tensor operations for better performance
        
        # Ensure inputs are contiguous for better memory access
        A = A.contiguous()
        B = B.contiguous()
        
        # Use bmm if batch dimensions exist, otherwise use mm
        if A.dim() == 3 and B.dim() == 3:
            return torch.bmm(A, B)
        else:
            return torch.mm(A, B)