Fixed @maxmenten's custom code to behave with new changes.