small fix: load at the module level, not in the sub function.

I merge it myself as this is a small fix that break the cuda code. Don't hesitate to review it.