+++++++++++++++++++++++++++++++++++++++++++++++++++++++
                      PAGE 4/5

+++++++++++++++++++++++++++++++++++++++++++++++++++++++
Code 4: Different deterministic randomness for each thread (GPU device).
In this case, each thread will generate different randomness using its own seed without affecting the other threads PRNGs, thanks to the locks.
import threading
import random

import numpy as np

import torch
import torch.nn as nn
from torch.nn import DataParallel


def set_seed(seed):
    """
    For seed to some modules.
    :param seed: int. The seed.
    :return:
    """
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        self.thread_lock = threading.Lock()

    def forward(self, x, seed):
        if x.is_cuda:
            device = x.get_device()
            print("DEVICE: {}".format(x.get_device()))
            print("x.size() = {}".format(x.size()))
            print("x:\n {}".format(x))
        else:
            device = torch.device("cpu")

        prngs0 = torch.random.get_rng_state().type(torch.float).numpy()
        # computing sum(abs(diff)**2) is a better way to have a summary of the PRNGs instead of computing only the sum.
        # Two different PRNGs may have the same sum (may happen if we generate few random numbers).
        print("DEVICE {} PRNG STATUS BEFORE: \n {}".format(torch.cuda.current_device(),
                                                           np.abs(np.diff(prngs0)**2).sum()))

        self.thread_lock.acquire()
        set_seed(seed)  # re-seed
        delta = torch.rand(x.size()).to(device)  # the danger zone.
        set_seed(seed)  # re-seed like nothing happened!!!
        self.thread_lock.release()

        prngs1 = torch.random.get_rng_state().type(torch.float).numpy()
        print("DEVICE {} PRNG STATUS AFTER: \n {}".format(torch.cuda.current_device(),
                                                          np.abs(np.diff(prngs1)**2).sum()))

        print("DEVICE {} DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): \n {}".format(
            torch.cuda.current_device(),   np.abs(prngs0 - prngs1).sum()))
        print("Delta:\n {}".format(delta))
        x = x + delta
        return x


if __name__ == "__main__":
    set_seed(0)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = Model()
    model = DataParallel(model)
    model.to(device)

    x = torch.rand(8, 3)
    x = x.to(device)
    print("X in:\n {}".format(x))
    print("X out:\n {}".format(model(x=x, seed=torch.tensor([10, 20]).to(device))))


run:
X in:
 tensor([[0.4963, 0.7682, 0.0885],
        [0.1320, 0.3074, 0.6341],
        [0.4901, 0.8964, 0.4556],
        [0.6323, 0.3489, 0.4017],
        [0.0223, 0.1689, 0.2939],
        [0.5185, 0.6977, 0.8000],
        [0.1610, 0.2823, 0.6816],
        [0.9152, 0.3971, 0.8742]], device='cuda:0')
DEVICE: 0
x.size() = torch.Size([4, 3])
DEVICE: 1
x.size() = torch.Size([4, 3])
x:
 tensor([[0.0223, 0.1689, 0.2939],
        [0.5185, 0.6977, 0.8000],
        [0.1610, 0.2823, 0.6816],
        [0.9152, 0.3971, 0.8742]], device='cuda:1')
DEVICE 1 PRNG STATUS BEFORE: 
 46790328.0
DEVICE 1 PRNG STATUS AFTER: 
 49251776.0
DEVICE 1 DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): 
 217190.0
Delta:
 tensor([[0.5615, 0.1774, 0.8147],
        [0.3295, 0.2319, 0.7832],
        [0.8544, 0.1012, 0.1877],
        [0.9310, 0.0899, 0.3156]], device='cuda:1')
x:
 tensor([[0.4963, 0.7682, 0.0885],
        [0.1320, 0.3074, 0.6341],
        [0.4901, 0.8964, 0.4556],
        [0.6323, 0.3489, 0.4017]], device='cuda:0')
DEVICE 0 PRNG STATUS BEFORE: 
 49251776.0
DEVICE 0 PRNG STATUS AFTER: 
 47405664.0
DEVICE 0 DIFF PRNG STATUS ABS(BEFORE - AFTER).SUM(): 
 218221.0
Delta:
 tensor([[0.4581, 0.4829, 0.3125],
        [0.6150, 0.2139, 0.4118],
        [0.6938, 0.9693, 0.6178],
        [0.3304, 0.5479, 0.4440]], device='cuda:0')
X out:
 tensor([[0.9543, 1.2511, 0.4010],
        [0.7471, 0.5214, 1.0459],
        [1.1839, 1.8658, 1.0734],
        [0.9627, 0.8968, 0.8457],
        [0.5838, 0.3463, 1.1086],
        [0.8481, 0.9295, 1.5832],
        [1.0154, 0.3834, 0.8693],
        [1.8462, 0.4870, 1.1898]], device='cuda:0')

The above code is deterministic, and produces different randomness per device.
Now, using code 4, one can produce reproducible results across a fixed number of devices K.