Function create_param_group 

pub fn create_param_group(
    params: Vec<PyTensor>,
    lr: f32,
    extra_params: HashMap<String, Py<PyAny>>,
) -> PyResult<HashMap<String, Py<PyAny>>>

Helper function to create parameter group