I wanted to follow up on this as I am looking rsample from a truncated Gaussian in PyTorch and compute log_prob and wanted to see if there were any updated developments. @anton’s module that he shared is very helpful, but unfortunately, I am looking for a solution that is CUDA-capable. Thank you for your help.