You can estimate the memory footprint of the model itself by summing the number of parameters, buffers (, and other tensors, if needed) and multiply it by the dtype factor (e.g. 4 for float32).
However, this would not give you the “complete” memory usage, since the forward activations (intermediates) as well as the gradients would also use memory. That being said, if you are using cudnn, then note that different algorithms would also consume a different amount of memory for their workspace, especially if you are using torch.backends.cudnn.benchmark = True.
I would thus recommend to perform an example training step with the shapes you are planning to use and check the memory usage e.g. via torch.cuda.memory_summary().