In the following line you are missing the mode parameter to kernel_product.
y_pred = relu(kernel_product(w1, x, s)).mm(w2)

This would be better
y_pred = relu(kernel_product(w1, x, "gaussian", s)).mm(w2)