Sorry for the late response.
The code you provide is the pseudo code I would give and very similar to the code I have in my code base, one small change is that the weights are generated by different samples.
def forward(self, samples, reference_samples):
    outputs = []
    for sample in samples:
        weight = generate_per_sample_weight(reference_samples)
        replace_conv_weight(self.conv, weight)
        outputs.append(self.conv(sample))
    return outputs

However, the main issue with this is that pytorch doesn’t distribute this over multiple GPU’s because of the for loop. These calculation are all located on the first main GPU (with the standard dataparallel package)
I was wondering if there is an easy way to write something like this which still allows for data parallelisation:
weights: Batch x # INPUT FILTERS x # OUTPUT FILTERS x FILTER WIDTH x FILTER HEIGHT
samples: BATCH x CHANNELS x WIDTH x HEIGHT
def forward(self, samples, weights):
    outputs = self.conv(sample, weights)
    return outputs

So this self.conv function would then be one purely based on matrices like the original conv one, which should allow data parallelisation.