hi Ptrblck,
If you use this code, then setting num workers to 0 uses 2891Mb and num workers=6 uses 6597Mb. it seems to be using the GPU memory in accordance with the following:
GPU_mem=3.8333num_work^3 - 78.5num_work^2 + 950.67*num_work + 2891
import itertools
ngpu=1
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

class dataset3D(Dataset):
    def __init__(self,transform=None):
        self.transform = transform
    def __len__(self):
        count=1000
        return count

    def __getitem__(self, index):

        data3dt=np.random.randint(0,1,(1,128,128,128))
        rescale_size = 99

        if np.max(data3dt)!=0:
            scrap = np.argwhere(data3dt == 1)

            max_r=np.max(scrap[:,0])
            min_r=np.min(scrap[:,0])
            max_c = np.max(scrap[:, 1])
            min_c=np.min(scrap[:,1])
            max_d = np.max(scrap[:, 2])
            min_d=np.min(scrap[:,2])
            resize3d_t = np.zeros((max_r+1, max_c+1, max_d+1))
            # print(resize3d_t.shape)
            r_range=max_r-min_r
            c_range = max_c - min_c
            d_range = max_d - min_d

            r_offset = min_r-int((resize3d_t.shape[0]-(max_r - min_r)) / 2)-1
            c_offset = min_c-int((resize3d_t.shape[1]-(max_c - min_c)) / 2)-1
            d_offset = min_d-int((resize3d_t.shape[2]-(max_d - min_d)) / 2)-1

            for count in range (len(scrap)):
                row=int(scrap[count,0]-r_offset)
                col=int(scrap[count,1]-c_offset)
                dep=int(scrap[count,2]-d_offset)
                resize3d_t[row,col,dep]=1
            for count in range (len(scrap)):
                row=int(scrap[count,0])
                col=int(scrap[count,1])
                dep=int(scrap[count,2])
                resize3d_t[row,col,dep]=1
            d3shape=resize3d_t.shape
            r_sf=d3shape[0]/rescale_size
            c_sf=d3shape[1]/rescale_size
            d_sf = d3shape[2] / rescale_size
            resize3d = np.zeros((rescale_size,rescale_size,rescale_size))
            for x, y, z in itertools.product(range(rescale_size-1),
                                             range(rescale_size-1),
                                             range(rescale_size-1)):
                resize3d[x][y][z] = resize3d_t[int(x * r_sf)][int(y * c_sf)][int(z * d_sf)]
            resize3d = torch.from_numpy(resize3d)
            resize3d = resize3d.unsqueeze(0)
            # print(resize3d.shape)
        else:
            resize3d=np.zeros((1,rescale_size,rescale_size,rescale_size))
            resize3d = torch.from_numpy(resize3d)
        data3d = {'resize': resize3d}
        if self.transform:
            data3d=self.transform(data3d)

        return data3d
class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, data3d):

        resize_array = data3d['resize']
        return {'resize':torch.from_numpy(resize_array)}

class D3_Encoder(nn.Module):
    def __init__(self, lat_dim):
        super(D3_Encoder, self).__init__()
        self.conv3d1 = nn.Conv3d(1, 96, (9,9,9), stride=(5,5,5), padding=(0, 0, 0))
        self.conv3d2 = nn.Conv3d(96,128, (9,9,9), stride=(5,5,5), padding=(0, 0, 0))

        self.fc1 = nn.Linear(3456, 256)
        self.fc2 = nn.Linear(256, lat_dim + lat_dim)

    def forward(self, data):
        data = self.conv3d1(data)

        data = self.conv3d2(data)

        data = data.view(-1, 3456)

        data = self.fc1(data)
        data = self.fc2(data)

        return data
class D3_Decoder(nn.Module):
    def __init__(self, lat_dim):
        super(D3_Decoder, self).__init__()
        self.conv3dT5 = nn.ConvTranspose3d(128, 96,(9,9,9), stride=(5,5,5), padding=(0, 0, 0))
        self.conv3dT6 = nn.ConvTranspose3d(96 , 1,(9,9,9), stride=(5,5,5), padding=(0, 0, 0))
        self.fc1 = nn.Linear(lat_dim, 256)
        self.fc2 = nn.Linear(256, 3456)

    def forward(self, z2):
        ###insert fully connected in here
        z2 = self.fc1(z2)
        z2 = self.fc2(z2)
        ### reshape the data before leaving the fully conected layer
        z2 = z2.view(-1, 128, 3, 3, 3)
        z2 = self.conv3dT5(z2)
        out = self.conv3dT6(z2)
        return out

class Gen_model_3D(nn.Module):
    def __init__(self, lat_dim=100):
        super(Gen_model_3D, self).__init__()
        self.D3_Encoder = D3_Encoder(lat_dim)
        self.D3_Decoder = D3_Decoder(lat_dim)
        self.cuda()
        self.lat_dim = lat_dim

    def D3_encode(self, data):
        data = model_3d.D3_Encoder(data)
        mean, logvar = torch.chunk(data, 2, dim=1)
        return mean, logvar

    def D3_decode(self, z2, apply_sigmoid=False):
        logits = self.D3_Decoder(z2)
        if apply_sigmoid:
            probs = torch.sigmoid(logits)
            return probs
        return logits

    def reparm(self, mean, logvar):
        epsilon = torch.randn(mean.size()).to(device)
        z = epsilon * torch.exp(logvar.mul(0.5)) + mean
        # print("z=",z.shape)
        return z
dataload=dataset3D(transform=None)

dataloader = torch.utils.data.DataLoader(dataload, batch_size=20,
                                         shuffle=True, num_workers=0)
def loss_fn(model_3d, data):
    mean, logvar = model_3d.D3_encode(data)
    z2=model_3d.reparm(mean, logvar)
    out=model_3d.D3_decode(z2)
    criterion = torch.nn.BCEWithLogitsLoss(size_average=False, reduce=True, reduction='none')
    BCE = criterion(out, data)
    KLD = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp().pow(2))
    return BCE + 3 * KLD, out

model_3d= Gen_model_3D()
optimizer_gen = torch.optim.Adam(model_3d.parameters(), lr=1E-5,
               weight_decay=1e-6)

def train():
    img_list=[]
    iters=0
    for epoch in range(10):
        for i,data in enumerate(dataloader):
            data_in=data['resize']

            data_in=data_in.cuda()
            data_in = data_in.type(torch.cuda.FloatTensor)
            # data_in = data_in.unsqueeze(1)

            optimizer_gen.zero_grad()
            loss,out = loss_fn(model_3d, data_in)
            loss.backward()
            print(loss)
            optimizer_gen.step()
if __name__ == '__main__':
    img_list =  train()