forked from ccf-ai-infra/GPUCodeForces
20 lines
633 B
Python
20 lines
633 B
Python
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) |