Modified the hardcoded bullshit

Fixed the hardcoded shit to deal with the dimension change.