Make dropout mask conditional on prob > 0