Hi,
I have two naive suggestions.

It seems like your feature_extraction0 and feature_extraction0 are the same module. Why not construct one and forward the combined input of x0 and x1?

Nvidia APEX 1 is a nice tool to speed up your training by leveraging 16-bit operations on the float tensor (instead of the default 32-bit)

I think it’s more clear if you can describe what the ResUNetBN2C and x0, x1 are, there may be methods to speed up depends on your case.