A major limiting factor in graphical model inference is the complexity of computing the partition function. Exact message-passing algorithms such as Bucket Elimination (BE) require exponentially high levels of memory to compute the partition function, therefore approximations are necessary. In this paper, we build upon a recently introduced methodology called Deep Bucket Elimination (DBE) that uses classical Neural Networks (NNs) to approximate messages generated by BE when buckets have large memory requirements. The main feature of our new scheme called NeuroBE is that it customizes the architecture and learning of the NNs to the message size and its distribution. We also explore a new loss function for training taking into account the estimated message cost distribution. Our experiments demonstrate that these enhancements provide significant improvements over DBE in both time and accuracy. We also study the impact of the messages local errors on the global accuracy of the estimate of the partition function.