So your saying that perhaps the size of the tensor may be too large to fit into a single memory “block”? (i dont know much about how memory is allocated), and thus takes up more memory than originally thought?
Any chance you could look at the network to see? I have it linked here: https://github.com/penguinshin/capsule-net/blob/master/capsule.py 3