One way to do this is to run the whole thing twice, as you mentioned, once in updateOutput and once in updateGradInput.
There is a way to share data between the two functions. You can define a python module that has forward() and backward() functions. The forward() function can all into the THNN updateOutput, and then save the result to a context. backward() is able to access this context, and you wouldn’t need a THNN updateGradInput in this case.
See here 4 for more details about this.
The more error checking, the better. Errors that occur in TH that aren’t checked very well can result in some hard to parse error messages (or problems) for the end user.