Update, I’ve got the following code working:
	def gated_branch_function(self, j_batch, control_batch):
		batch_outputs = []
		for j, control in zip(j_batch, control_batch):
			control = int(control.item())
			""" Branches """
			if control == 2 or control == 0:
				output = self.follow_lane_branch(j)
			elif control == 3:
				output = self.left_branch(j)
			elif control == 4:
				output = self.right_branch(j)
			elif control == 5:
				output = self.straight_branch(j)
			else:
				output = self.general_branch(j)
			batch_outputs.append(output)
		return torch.stack(batch_outputs)


	def forward(self, input_data):
		''' Define variables '''
		input_image = input_data[0].permute(0,3,2,1).float()
		input_speed = input_data[1].unsqueeze(dim=1)
		control = input_data[2]

		''' Pass input data into network '''
		imageOutput = self.imageModule(input_image)
		speedOutput = self.measurementModule(input_speed)

		''' Joint sensory '''
		j = torch.cat([imageOutput, speedOutput], 1)
		j = self.jointSensoryModule(j)

		''' Branches '''
		output = self.gated_branch_function(j, control)
		return output

However this results in a bottleneck, so training is slow. Is there anyway I could write gated_branch_function so that each sample in the batch gets sent to their respective branch in parallel?