Add missing task types to the gathering of weights in the partition code

Should be a trivial and harmless change (and has likely no impact) but it's probably better to have all the task types that actually do substantial work here.