Allow overriding the batch size on a per dataset level

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1232,6 +1232,17 @@ export default function SimpleJob({
                         placeholder="eg. 1"
                         docKey={'dataset.num_repeats'}
                       />
+                      <NumberInput
+                        label="Batch Size"
+                        value={dataset.batch_size ?? null}
+                        className="pt-2"
+                        onChange={value =>
+                          setJobConfig(value == null ? undefined : value, `config.process[0].datasets[${i}].batch_size`)
+                        }
+                        placeholder={`${jobConfig.config.process[0].train.batch_size}`}
+                        min={1}
+                        allowEmpty
+                      />
                     </div>
                     <div>
                       <TextInput

ui/src/components/formInputs.tsx

@@ -152,10 +152,12 @@ export interface NumberInputProps extends InputProps {
   onChange: (value: number | null) => void;
   min?: number;
   max?: number;
+  // when true, clearing the input calls onChange(null) instead of being ignored
+  allowEmpty?: boolean;
 }
 
 export const NumberInput = (props: NumberInputProps) => {
-  const { label, value, onChange, placeholder, required, min, max, docKey = null } = props;
+  const { label, value, onChange, placeholder, required, min, max, allowEmpty, docKey = null } = props;
   let { doc } = props;
   if (!doc && docKey) {
     doc = getDoc(docKey);
@@ -193,6 +195,9 @@ export const NumberInput = (props: NumberInputProps) => {
           // Handle empty or partial inputs
           if (rawValue === '' || rawValue === '-') {
             // For empty or partial negative input, don't call onChange yet
+            if (rawValue === '' && allowEmpty) {
+              onChange(null);
+            }
             return;
           }
 

ui/src/types.ts

@@ -106,6 +106,7 @@ export interface SaveConfig {
 }
 
 export interface DatasetConfig {
+  batch_size?: number;
   folder_path: string;
   mask_path: string | null;
   mask_min_value: number;