Repairing cast bf16/f16

1 files changed, 4 insertions, 4 deletions

diff --git a/candle-kernels/src/cast.cu b/candle-kernels/src/cast.cu
index ea611eba..0c4ddbc6 100644
--- a/candle-kernels/src/cast.cu
+++ b/candle-kernels/src/cast.cu
@@ -13,13 +13,13 @@ extern "C" __global__ void FN_NAME( \
     const size_t *strides = info + num_dims; \
     if (is_contiguous(num_dims, dims, strides)) { \
         for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
-            out[i] = inp[i]; \
+            out[i] = (DST_TYPENAME) inp[i]; \
         } \
     } \
     else { \
         for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
             unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \
-            out[i] = inp[strided_i]; \
+            out[i] = (DST_TYPENAME) inp[strided_i]; \
         } \
     } \
 } \
@@ -29,12 +29,12 @@ CAST_OP(__nv_bfloat16, __nv_bfloat16, cast_bf16_bf16)
 
 // CAST_OP(__nv_bfloat16, uint8_t, cast_bf16_u8)
 CAST_OP(__nv_bfloat16, uint32_t, cast_bf16_u32)
-// CAST_OP(__nv_bfloat16, __half,   cast_bf16_f16)
+CAST_OP(__nv_bfloat16, __half,   cast_bf16_f16)
 CAST_OP(__nv_bfloat16, float,    cast_bf16_f32)
 CAST_OP(__nv_bfloat16, double,   cast_bf16_f64)
 CAST_OP(uint8_t, __nv_bfloat16, cast_u8_bf16)
 CAST_OP(uint32_t, __nv_bfloat16, cast_u32_bf16)
-// CAST_OP(__half,   __nv_bfloat16, cast_f16_bf16)
+CAST_OP(__half,   __nv_bfloat16, cast_f16_bf16)
 CAST_OP(float,    __nv_bfloat16, cast_f32_bf16)
 CAST_OP(double,   __nv_bfloat16, cast_f64_bf16)
 #endif