← Back to Leaderboard

The AI CUDA Engineer 👷

76_Gemm_Add_ReLU • combined_warp_tile_base

Level 2 • Task 76
import torch
import torch.nn as nn
import torch.nn.functional as F


def module_fn(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
) -> torch.Tensor:
    """
    Performs matrix multiplication, adds bias, and applies ReLU activation.

    Args:
        x (torch.Tensor): Input tensor with shape (batch_size, in_features)
        weight (torch.Tensor): Weight matrix with shape (out_features, in_features)
        bias (torch.Tensor): Bias tensor with shape (out_features,)

    Returns:
        torch.Tensor: Output tensor with shape (batch_size, out_features)
    """
    x = F.linear(x, weight)
    x = x + bias
    x = F.relu(x)
    return x


class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, adds a bias term, and applies ReLU.
    """

    def __init__(self, in_features, out_features, bias_shape):
        super(Model, self).__init__()
        gemm = nn.Linear(in_features, out_features, bias=False)
        self.weight = nn.Parameter(gemm.weight)
        self.bias = nn.Parameter(torch.randn(bias_shape) * 0.02)

    def forward(self, x, fn=module_fn):
        return fn(x, self.weight, self.bias)


batch_size = 128
in_features = 1024
out_features = 512
bias_shape = (out_features,)


def get_inputs():
    return [torch.randn(batch_size, in_features)]


def get_init_inputs():
    return [in_features, out_features, bias_shape]

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Simple model that performs a matrix multiplication, adds a bias term, and applies ReLU.
    """
    def __init__(self, in_features, out_features, bias_shape):
        super(Model, self).__init__()
        self.gemm = nn.Linear(in_features, out_features, bias=False)
        self.bias = nn.Parameter(torch.randn(bias_shape)*0.02)

    def forward(self, x):   
        """
        Args:
            x (torch.Tensor): Input tensor with shape (batch_size, in_features).
        Returns:
            torch.Tensor: Output tensor with shape (batch_size, out_features).
        """
        x = self.gemm(x)
        x = x + self.bias
        x = torch.relu(x)
        return x

batch_size = 128
in_features = 1024
out_features = 512
bias_shape = (out_features,)

def get_inputs():
    return [torch.randn(batch_size, in_features)]

def get_init_inputs():
    return [in_features, out_features, bias_shape]
Download Evaluation Download PyTorch Download CUDA Download Profiles

Kernel Information

Operation Name 76_Gemm_Add_ReLU
Level ID 2
Task ID 76
Kernel Name combined_warp_tile_base
CUDA Speedup (Native) 0.926x
CUDA Speedup (Compile) 1.542x
CUDA Runtime 0.029 ms
PyTorch Runtime (Native) 0.027 ms
PyTorch Runtime (Compile) 0.045 ms
Correct True
Max Diff (vs. Reference) 0.000000
Model o3-mini-2025-01-31
Temperature 1.00
View Experiment Progress Details

Related Kernels (Level 2, Task 76 • 76_Gemm_Add_ReLU)

Rank Kernel Name Runtime (ms) Speedup Native Speedup Compile
🥇 shared_warp_tile_kernel_base 0.03 0.93 1.54
🥇 combined_warp_tile_base 0.03 0.93 1.54
🥉 optimized_block_size_kernel_base 0.03 0.89 1.49
4 warp_tile_ldg_base 0.03 0.87 1.44
4 even_workload_dist_base_base 0.03 0.87 1.44
4 hybrid_warp_tile_kernel_base 0.03 0.87 1.44
4 warp_tile_hybrid_base 0.03 0.87 1.44
8 warp_tile_ldg_opt_base 0.03 0.81 1.36
8 warp_reduction_optimized_base_base 0.03 0.81 1.36
10 optimized_shared_memory_base_base 0.03 0.79 1.32
10 warp_tile_base_base 0.03 0.79 1.32
12 hybrid_optimized_kernel_base 0.04 0.77 1.28
13 warp_reduction_gemm_base 0.04 0.71 1.18
13 warp_tile_aligned_base_base 0.04 0.71 1.18
15 vectorized_warp_unroll_base_base 0.04 0.69 1.15
15 vectorized_warp_unroll_base_edit_1 0.04 0.69 1.15
15 warp_reduction_unrolled_gemm_edit_1 0.04 0.69 1.15
18 unrolled_warp_gemm_edit_1 0.04 0.67 1.12
18 unrolled_warp_gemm_base 0.04 0.67 1.12
18 vectorized_warp_reduction_base 0.04 0.67 1.12 
/*
Combined CUDA kernel that fuses warp-level tiling with vectorized memory accesses and loop unrolling
from two different implementations. Each warp processes a tile of output features (TILE_SIZE outputs)
for a given batch sample. The kernel uses __ldg() for read-only loads and vectorized float4 loads
for 128-bit aligned accesses. Remaining elements are unrolled as in the second kernel. 
*/

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAGuard.h>

#define WARP_SIZE 32
#define TILE_SIZE 4  // Each warp processes TILE_SIZE output elements

// Combined kernel: each block handles one or more batch samples along grid.x, and groups of output features along grid.y.
// In each block, warps are assigned a contiguous tile of TILE_SIZE outputs. Within each tile, threads cooperatively
// compute the dot product using vectorized loads and handle any remainder elements via loop unrolling.

__global__ void combined_warp_tile_kernel(const float* __restrict__ x,
                                           const float* __restrict__ weight,
                                           const float* __restrict__ bias,
                                           float* __restrict__ out,
                                           int in_features,
                                           int out_features) {
    // Each block processes one batch sample (grid.x) and a group of output features (grid.y).
    int batch_idx = blockIdx.x;

    // Calculate warp and lane indices
    int warps_per_block = blockDim.x / WARP_SIZE;  // e.g., 8 warps per block
    int warp_id = threadIdx.x / WARP_SIZE;           
    int lane_id = threadIdx.x % WARP_SIZE;

    // Compute the starting output index for this warp's tile
    int base_out_group = blockIdx.y * (warps_per_block * TILE_SIZE);
    int out_base = base_out_group + warp_id * TILE_SIZE;

    // Early exit if the starting output index is out-of-bounds
    if (out_base >= out_features) return;

    // Pointer to the current batch row
    const float* x_row = x + batch_idx * in_features;

    // Array to hold partial sums for each output within the tile
    float sums[TILE_SIZE] = {0.0f, 0.0f, 0.0f, 0.0f};

    // Determine vectorized loop parameters
    int nvec = in_features / 4;  // number of float4 loads
    int rem = in_features % 4;   // remaining elements

    // Process each output feature in the tile
    #pragma unroll
    for (int tile = 0; tile < TILE_SIZE; tile++) {
        int current_out = out_base + tile;
        if (current_out < out_features) {
            const float* w_row = weight + current_out * in_features;

            // Cast pointers for vectorized loads (assumes data is 128-bit aligned)
            const float4* x_vec = reinterpret_cast<const float4*>(x_row);
            const float4* w_vec = reinterpret_cast<const float4*>(w_row);

            float sum = 0.0f;

            // Main vectorized loop: each thread handles a stride of WARP_SIZE
            for (int k = lane_id; k < nvec; k += WARP_SIZE) {
                float4 a = __ldg(&x_vec[k]);
                float4 b = __ldg(&w_vec[k]);
                sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
            }

            // Handle remainder elements (if any) with loop unrolling
            int offset = nvec * 4;
            for (int r = lane_id; r < rem; r += WARP_SIZE) {
                sum += __ldg(x_row + offset + r) * __ldg(w_row + offset + r);
            }

            sums[tile] = sum;
        }
    }

    // Perform warp-level reduction using shuffle instructions
    #pragma unroll
    for (int tile = 0; tile < TILE_SIZE; tile++) {
        for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
            sums[tile] += __shfl_down_sync(0xffffffff, sums[tile], offset);
        }
    }

    // The first lane of each warp writes the final output with bias and ReLU activation
    if (lane_id == 0) {
        for (int tile = 0; tile < TILE_SIZE; tile++) {
            int current_out = out_base + tile;
            if (current_out < out_features) {
                float result = sums[tile] + __ldg(bias + current_out);
                out[batch_idx * out_features + current_out] = (result > 0.0f) ? result : 0.0f;
            }
        }
    }
}

// Host function to launch the combined kernel
// Grid configuration: blockIdx.x corresponds to batch index; blockIdx.y covers groups of output features.
// Block configuration: fixed number of threads, with each warp handling TILE_SIZE outputs.

torch::Tensor combined_linear_relu_forward(torch::Tensor x,
                                             torch::Tensor weight,
                                             torch::Tensor bias) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");

    int batch_size = x.size(0);
    int in_features = x.size(1);
    int out_features = weight.size(0);

    auto out = torch::empty({batch_size, out_features}, x.options());

    // Configure execution parameters
    int warps_per_block = 8;  // can be tuned
    int threads_per_block = warps_per_block * WARP_SIZE;
    int blocks_y = (out_features + (warps_per_block * TILE_SIZE) - 1) / (warps_per_block * TILE_SIZE);

    dim3 grid(batch_size, blocks_y);
    dim3 block(threads_per_block);

    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    combined_warp_tile_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        in_features,
        out_features
    );

    return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &combined_linear_relu_forward, "Combined GEMM with bias and ReLU (CUDA) using warp-level tile and vectorized memory access");
}
Performance Metrics
Metric Value Unit Variance Samples
Executed Ipc Active 1.620 inst/cycle 0.000 5
Executed Ipc Elapsed 1.478 inst/cycle 0.000 5
Issue Slots Busy 40.642 % 0.022 5
Issued Ipc Active 1.626 inst/cycle 0.000 5
SM Busy 40.642 % 0.022 5
Memory Throughput 86557181097.354 byte/second 1641481626324375808.000 5
Mem Busy 76.324 % 1.167 5
Max Bandwidth 73.742 % 1.083 5
L1/TEX Hit Rate 55.848 % 0.265 5
L2 Hit Rate 90.586 % 6.879 5
Mem Pipes Busy 25.448 % 0.131 5
Warp Cycles Per Issued Instruction 25.256 cycle 0.001 5
Warp Cycles Per Executed Instruction 25.342 cycle 0.001 5
Avg. Active Threads Per Warp 30.470 0.000 5
Avg. Not Predicated Off Threads Per Warp 28.840 0.000 5
Max Active Clusters 0.000 cluster 0.000 5
Max Cluster Size 8.000 block 0.000 5
Overall GPU Occupancy 0.000 % 0.000 5
Cluster Occupancy 0.000 % 0.000 5
Block Limit SM 32.000 block 0.000 5
Block Limit Registers 6.000 block 0.000 5
Block Limit Shared Mem 32.000 block 0.000 5
Block Limit Warps 8.000 block 0.000 5
Theoretical Active Warps per SM 48.000 warp 0.000 5
Theoretical Occupancy 75.000 % 0.000 5
Achieved Occupancy 64.140 % 0.060 5
Achieved Active Warps Per SM 41.050 warp 0.024 5
Analysis Rules
Rule Description
INF HighPipeUtilization ALU is the highest-utilized pipeline (22.3%) based on active cycles, taking into account the rates of its different instructions. It executes integer and logic operations. It is well-utilized, but should not be a bottleneck.
INF CPIStall Check the Warp Stall Sampling (All Cycles) table for the top stall locations in your source based on sampling data. The Kernel Profiling Guide (https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#metrics-reference) provides more details on each stall reason.
WRN Occupancy This kernel's theoretical occupancy (75.0%) is limited by the number of required registers. The difference between calculated theoretical (75.0%) and measured achieved occupancy (63.9%) can be the result of warp scheduling overheads or workload imbalances during the kernel execution. Load imbalances can occur between warps within a block as well as across blocks of the same kernel. See the CUDA Best Practices Guide (https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/index.html#occupancy) for more details on optimizing occupancy.
Operation / Metric Value Unit
aten::to
CPU Time 666284.28 μs
Device Time 221.25 μs
Self CPU Time 70.55 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::_to_copy
CPU Time 666213.73 μs
Device Time 221.25 μs
Self CPU Time 120.94 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::empty_strided
CPU Time 665520.99 μs
Device Time 0.00 μs
Self CPU Time 148.61 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaDeviceGetStreamPriorityRange
CPU Time 649629.61 μs
Device Time 0.00 μs
Self CPU Time 649629.61 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
cudaLaunchKernel
CPU Time 324214.50 μs
Device Time 548.32 μs
Self CPU Time 324214.50 μs
Self Device Time 548.32 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
combined_warp_tile_kernel(float const*, float const*, float const*, float*, int, int)
CPU Time 0.00 μs
Device Time 138738.76 μs
Self CPU Time 0.00 μs
Self Device Time 138738.76 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::zero_
CPU Time 90906.42 μs
Device Time 413187.97 μs
Self CPU Time 14625.65 μs
Self Device Time 0.00 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
aten::fill_
CPU Time 76282.59 μs
Device Time 413187.97 μs
Self CPU Time 18702.27 μs
Self Device Time 413187.97 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
void at::native::vectorized_elementwise_kernel<4, at::native::FillFunctor<int>, at::detail::Array<char*, 1> >(int, at::native::FillFunctor<int>, at::detail::Array<char*, 1>)
CPU Time 0.00 μs
Device Time 413187.97 μs
Self CPU Time 0.00 μs
Self Device Time 413187.97 μs
CPU Memory Usage 0 B
Device Memory Usage 0 B
Self CPU Memory Usage 0 B
Self Device Memory Usage 0 B
Status: Completed 
45315 warnings generated when compiling for host.
Suppressed 45347 warnings (45300 in non-user code, 47 NOLINT).
Use -header-filter=.* to display errors from all non-system headers. Use -system-headers to display errors from system headers as well.
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:20:43 bugprone-easily-swappable-parameters
20 | __global__ void combined_warp_tile_kernel(const float* __restrict__ x,
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
21 | const float* __restrict__ weight,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
22 | const float* __restrict__ bias,
| ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:20:69: note: the first parameter in the range is 'x'
20 | __global__ void combined_warp_tile_kernel(const float* __restrict__ x,
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:22:70: note: the last parameter in the range is 'bias'
22 | const float* __restrict__ bias,
| ^~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:24:44: warning: 2 adjacent parameters of 'combined_warp_tile_kernel' of similar type ('int') are easily swapped by mistake [bugprone-easily-swappable-parameters]
24 | int in_features,
| ^~~~~~~~~~~~~~~~
25 | int out_features) {
| ~~~~~~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:24:48: note: the first parameter in the range is 'in_features'
24 | int in_features,
| ^~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:25:48: note: the last parameter in the range is 'out_features'
25 | int out_features) {
| ^~~~~~~~~~~~
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:27:21: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
27 | int batch_idx = blockIdx.x;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:30:27: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
30 | int warps_per_block = blockDim.x / WARP_SIZE; // e.g., 8 warps per block
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:31:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
31 | int warp_id = threadIdx.x / WARP_SIZE;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:32:19: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
32 | int lane_id = threadIdx.x % WARP_SIZE;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:35:26: warning: narrowing conversion from 'unsigned int' to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
35 | int base_out_group = blockIdx.y * (warps_per_block * TILE_SIZE);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:42:26: warning: result of multiplication in type 'int' is used as a pointer offset after an implicit widening conversion to type 'ptrdiff_t' [bugprone-implicit-widening-of-multiplication-result]
42 | const float* x_row = x + batch_idx * in_features;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:42:30: note: make conversion explicit to silence this warning
12 | const float* x_row = x + batch_idx * in_features;
| ^~~~~~~~~~~~~~~~~~~~~~~
| static_cast<ptrdiff_t>()
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:42:30: note: perform multiplication in a wider type
42 | const float* x_row = x + batch_idx * in_features;
| ^~~~~~~~~
| static_cast<ptrdiff_t>()
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:56:34: warning: result of multiplication in type 'int' is used as a pointer offset after an implicit widening conversion to type 'ptrdiff_t' [bugprone-implicit-widening-of-multiplication-result]
56 | const float* w_row = weight + current_out * in_features;
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:56:43: note: make conversion explicit to silence this warning
56 | const float* w_row = weight + current_out * in_features;
| ^~~~~~~~~~~~~~~~~~~~~~~~~
| static_cast<ptrdiff_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:56:43: note: perform multiplication in a wider type
56 | const float* w_row = weight + current_out * in_features;
| ^~~~~~~~~~~
| static_cast<ptrdiff_t>( )
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:105:58: warning: the parameter 'x' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
105 | torch::Tensor combined_linear_relu_forward(torch::Tensor x,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:106:60: warning: the parameter 'weight' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
106 | torch::Tensor weight,
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:107:60: warning: the parameter 'bias' is copied for each invocation but only used as a const reference; consider making it a const reference [performance-unnecessary-value-param]
107 | torch::Tensor bias) {
| ^
| const &
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:112:22: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
112 | int batch_size = x.size(0);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:113:23: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
113 | int in_features = x.size(1);
| ^
/home/robert_sakana_ai/llm_cuda/experiments/20250203_optimize_b10_s4_e0_sweep/level_2/task_76/b8_s3_combined_warp_tile/base/base.cu:114:24: warning: narrowing conversion from 'int64_t' (aka 'long') to signed type 'int' is implementation-defined [bugprone-narrowing-conversions]
114 | int out_features = weight.size(0);
| ^