Skip to content

Commit f1737de

Browse files
JohannesGaesslerNexesenex
authored andcommitted
CUDA: fix 1D im2col, add tests (ggml/993)
1 parent 05969d1 commit f1737de

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

ggml/src/ggml-cuda/im2col.cu

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,9 @@ void ggml_cuda_op_im2col(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
9191
const int64_t OH = is_2D ? dst->ne[2] : 1;
9292
const int64_t OW = dst->ne[1];
9393

94-
const size_t delta_offset = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
95-
const int64_t batch = src1->ne[3];
96-
const size_t batch_offset = src1->nb[3] / 4; // nb is byte offset, src is type float32
94+
const size_t delta_offset = src1->nb[is_2D ? 2 : 1] / 4; // nb is byte offset, src is type float32
95+
const int64_t batch = src1->ne[is_2D ? 3 : 2];
96+
const size_t batch_offset = src1->nb[is_2D ? 3 : 2] / 4; // nb is byte offset, src is type float32
9797

9898
if(dst->type == GGML_TYPE_F16) {
9999
im2col_cuda_f16(src1_d, (half *) dst_d, IW, IH, OW, OH, KW, KH, IC, batch, batch_offset, delta_offset, s0, s1, p0, p1, d0, d1, stream);

0 commit comments

Comments
 (0)